diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 86fb33a82ce04..b145ab89b6962 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -222,6 +222,33 @@ jobs: # - name: Check datafusion (no-default-features) run: cargo xtask ci step check datafusion no-default + - name: Check datafusion (object_store) + run: cargo xtask ci step check datafusion object_store + - name: Check datafusion does not require object_store + run: | + tree_file=$(mktemp) + trap 'rm -f "$tree_file"' EXIT + cargo tree -p datafusion --no-default-features --features sql \ + --target all --edges normal --prefix none --format '{p}' > "$tree_file" + if grep -q '^object_store v' "$tree_file"; then + cat "$tree_file" + echo "DataFusion without storage must not depend on object_store" >&2 + exit 1 + fi + - name: Test file storage without object_store + run: | + cargo test -p datafusion --no-default-features \ + --features sql,parquet,avro,opendal,parquet_encryption --test storage + tree_file=$(mktemp) + trap 'rm -f "$tree_file"' EXIT + cargo tree -p datafusion --no-default-features \ + --features sql,parquet,avro,opendal,parquet_encryption \ + --edges normal --prefix none --format '{p}' > "$tree_file" + if grep -q '^object_store v' "$tree_file"; then + cat "$tree_file" + echo "File storage must not depend on object_store" >&2 + exit 1 + fi - name: Check datafusion (nested_expressions) run: cargo xtask ci step check datafusion nested_expressions - name: Check datafusion (array_expressions) diff --git a/Cargo.lock b/Cargo.lock index 18fde39a95924..bdfa59902d8f2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -399,7 +399,7 @@ version = "59.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "10fab8d4563491417ba801fab29d205104d20d4bdf37bda6cd1cf425cff598cd" dependencies = [ - "bitflags", + "bitflags 2.11.1", "serde", "serde_core", "serde_json", @@ -514,6 +514,16 @@ dependencies = [ "syn 3.0.5", ] +[[package]] +name = "asyncband" +version = "0.6.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94a214ba60d6231afd0e805e3c27c45a1626d9debaa5a5061c45a1ea1b2f1ed0" +dependencies = [ + "hashbrown 0.17.1", + "slab", +] + [[package]] name = "atoi" version = "2.0.0" @@ -1015,6 +1025,12 @@ dependencies = [ "num-traits", ] +[[package]] +name = "bitflags" +version = "1.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" + [[package]] name = "bitflags" version = "2.11.1" @@ -1069,7 +1085,7 @@ checksum = "ee04c4c84f1f811b017f2fbb7dd8815c976e7ca98593de9c1e2afad0f636bff4" dependencies = [ "async-stream", "base64 0.22.1", - "bitflags", + "bitflags 2.11.1", "bollard-buildkit-proto", "bollard-stubs", "bytes", @@ -1740,6 +1756,9 @@ dependencies = [ "datafusion-physical-plan", "datafusion-session", "datafusion-sql", + "datafusion-storage", + "datafusion-storage-object-store", + "datafusion-storage-opendal", "doc-comment", "env_logger", "flate2", @@ -1753,6 +1772,7 @@ dependencies = [ "log", "nix", "object_store", + "opendal", "parking_lot", "parquet", "rand 0.9.4", @@ -1786,6 +1806,7 @@ dependencies = [ "datafusion-common", "datafusion-common-runtime", "datafusion-proto", + "datafusion-storage-object-store", "env_logger", "futures", "libmimalloc-sys", @@ -1819,6 +1840,7 @@ dependencies = [ "datafusion-physical-expr", "datafusion-physical-plan", "datafusion-session", + "datafusion-storage", "futures", "itertools 0.15.0", "log", @@ -1844,6 +1866,7 @@ dependencies = [ "datafusion-physical-expr-adapter", "datafusion-physical-expr-common", "datafusion-physical-plan", + "datafusion-storage", "futures", "itertools 0.15.0", "log", @@ -1865,6 +1888,7 @@ dependencies = [ "datafusion", "datafusion-common", "datafusion-spark", + "datafusion-storage-object-store", "dirs", "env_logger", "futures", @@ -1893,6 +1917,7 @@ dependencies = [ "arrow-schema", "chrono", "criterion", + "datafusion-storage", "foldhash 0.2.0", "half", "hashbrown 0.17.1", @@ -1903,7 +1928,6 @@ dependencies = [ "libc", "log", "num-traits", - "object_store", "parquet", "rand 0.9.4", "recursive", @@ -1944,6 +1968,7 @@ dependencies = [ "datafusion-physical-plan", "datafusion-proto-models", "datafusion-session", + "datafusion-storage", "flate2", "futures", "glob", @@ -1955,6 +1980,7 @@ dependencies = [ "parking_lot", "rand 0.9.4", "tempfile", + "test-utils", "tokio", "tokio-util", "url", @@ -1979,9 +2005,11 @@ dependencies = [ "datafusion-physical-plan", "datafusion-proto-models", "datafusion-session", + "datafusion-storage", "futures", "itertools 0.15.0", "object_store", + "test-utils", "tokio", ] @@ -1999,8 +2027,10 @@ dependencies = [ "datafusion-physical-plan", "datafusion-proto-models", "datafusion-session", + "datafusion-storage", "futures", "object_store", + "test-utils", ] [[package]] @@ -2019,9 +2049,11 @@ dependencies = [ "datafusion-physical-plan", "datafusion-proto-models", "datafusion-session", + "datafusion-storage", "futures", "object_store", "regex", + "test-utils", "tokio", ] @@ -2041,9 +2073,11 @@ dependencies = [ "datafusion-physical-plan", "datafusion-proto-models", "datafusion-session", + "datafusion-storage", "futures", "object_store", "serde_json", + "test-utils", "tokio", "tokio-stream", ] @@ -2073,6 +2107,7 @@ dependencies = [ "datafusion-proto-models", "datafusion-pruning", "datafusion-session", + "datafusion-storage", "futures", "itertools 0.15.0", "log", @@ -2080,7 +2115,9 @@ dependencies = [ "parking_lot", "parquet", "tempfile", + "test-utils", "tokio", + "url", ] [[package]] @@ -2104,6 +2141,7 @@ dependencies = [ "datafusion-physical-expr-adapter", "datafusion-proto", "datafusion-sql", + "datafusion-storage-object-store", "env_logger", "futures", "insta", @@ -2141,10 +2179,10 @@ dependencies = [ "datafusion-common", "datafusion-expr", "datafusion-physical-expr-common", + "datafusion-storage", "futures", "insta", "log", - "object_store", "parking_lot", "parquet", "pin-project-lite", @@ -2548,6 +2586,7 @@ dependencies = [ "datafusion-physical-plan", "datafusion-proto-common", "datafusion-proto-models", + "datafusion-storage", "doc-comment", "flate2", "object_store", @@ -2556,6 +2595,7 @@ dependencies = [ "recursive", "serde_json", "tempfile", + "test-utils", "tokio", ] @@ -2684,6 +2724,7 @@ dependencies = [ "clap", "datafusion", "datafusion-spark", + "datafusion-storage", "datafusion-substrait", "env_logger", "futures", @@ -2705,6 +2746,49 @@ dependencies = [ "tokio-postgres", ] +[[package]] +name = "datafusion-storage" +version = "55.0.0" +dependencies = [ + "async-trait", + "bytes", + "chrono", + "futures", + "parking_lot", + "percent-encoding", + "thiserror", + "tokio", + "tokio-util", + "url", +] + +[[package]] +name = "datafusion-storage-object-store" +version = "55.0.0" +dependencies = [ + "async-trait", + "bytes", + "chrono", + "datafusion-storage", + "futures", + "object_store", + "tokio", +] + +[[package]] +name = "datafusion-storage-opendal" +version = "55.0.0" +dependencies = [ + "async-trait", + "bytes", + "chrono", + "datafusion-storage", + "futures", + "opendal", + "tokio", + "tokio-util", +] + [[package]] name = "datafusion-substrait" version = "55.0.0" @@ -2714,6 +2798,7 @@ dependencies = [ "chrono", "datafusion", "datafusion-functions-aggregate", + "datafusion-storage", "half", "insta", "itertools 0.15.0", @@ -2740,6 +2825,7 @@ dependencies = [ "datafusion-optimizer", "datafusion-physical-plan", "datafusion-sql", + "datafusion-storage-object-store", "futures", "getrandom 0.3.4", "object_store", @@ -2756,6 +2842,37 @@ dependencies = [ "insta", ] +[[package]] +name = "defmt" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2953bfe4f93bbd20cc71198842756f77d161884c99ebbabc41d80231ded88d1" +dependencies = [ + "bitflags 1.3.2", + "defmt-macros", +] + +[[package]] +name = "defmt-macros" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bad9c72e7ca2137e0dc3813245a0d282fd6daad32fd800af018306a9169b5fe8" +dependencies = [ + "defmt-parser", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "defmt-parser" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10d60334b3b2e7c9d91ef8150abfb6fa4c1c39ebbcf4a81c2e346aad939fee3e" +dependencies = [ + "thiserror", +] + [[package]] name = "deranged" version = "0.5.8" @@ -2821,7 +2938,7 @@ version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e0e367e4e7da84520dedcac1901e4da967309406d1e51017ae1abfb97adbd38" dependencies = [ - "bitflags", + "bitflags 2.11.1", "objc2", ] @@ -3056,7 +3173,7 @@ version = "25.12.19" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "35f6839d7b3b98adde531effaf34f0c2badc6f4735d26fe74709d8e513a96ef3" dependencies = [ - "bitflags", + "bitflags 2.11.1", "rustc_version", ] @@ -3845,28 +3962,59 @@ checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" [[package]] name = "jiff" -version = "0.2.24" +version = "0.2.35" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f00b5dbd620d61dfdcb6007c9c1f6054ebd75319f163d886a9055cec1155073d" +checksum = "668b7183bd07af9a4885f5c35b0cc5c83c4607a913c16b7e17291832910d2dcc" dependencies = [ + "defmt", + "jiff-core", "jiff-static", + "jiff-tzdb-platform", + "js-sys", "log", "portable-atomic", "portable-atomic-util", "serde_core", + "wasm-bindgen", + "windows-link", +] + +[[package]] +name = "jiff-core" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7feca88439efe53da3754500c1851dedf3cb36c524dd5cf8225cc0794de95d09" +dependencies = [ + "defmt", ] [[package]] name = "jiff-static" -version = "0.2.24" +version = "0.2.35" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e000de030ff8022ea1da3f466fbb0f3a809f5e51ed31f6dd931c35181ad8e6d7" +checksum = "3a69dcb3a21cfb32ce1cd056169337ca284af0766dd766e7878819b251a49204" dependencies = [ + "jiff-core", "proc-macro2", "quote", "syn 2.0.119", ] +[[package]] +name = "jiff-tzdb" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "142bd39932ad231f10513df9ab62661fead8719872150b7ad02a2df79f4e141e" + +[[package]] +name = "jiff-tzdb-platform" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "875a5a69ac2bab1a891711cf5eccbec1ce0341ea805560dcd90b7a2e925132e8" +dependencies = [ + "jiff-tzdb", +] + [[package]] name = "jobserver" version = "0.1.34" @@ -4022,7 +4170,7 @@ version = "0.1.16" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e02f3bb43d335493c96bf3fd3a321600bf6bd07ed34bc64118e9293bdffea46c" dependencies = [ - "bitflags", + "bitflags 2.11.1", "libc", "plain", "redox_syscall 0.7.4", @@ -4193,7 +4341,7 @@ version = "0.31.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf20d2fde8ff38632c426f1165ed7436270b44f199fc55284c38276f9db47c3d" dependencies = [ - "bitflags", + "bitflags 2.11.1", "cfg-if", "cfg_aliases", "libc", @@ -4331,7 +4479,7 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2a180dd8642fa45cdb7dd721cd4c11b1cadd4929ce112ebd8b9f5803cc79d536" dependencies = [ - "bitflags", + "bitflags 2.11.1", "dispatch2", "objc2", ] @@ -4348,7 +4496,7 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3e0adef53c21f888deb4fa59fc59f7eb17404926ee8a6f59f5df0fd7f9f3272" dependencies = [ - "bitflags", + "bitflags 2.11.1", "objc2", ] @@ -4413,7 +4561,7 @@ dependencies = [ "md-5 0.10.6", "parking_lot", "percent-encoding", - "quick-xml", + "quick-xml 0.39.2", "rand 0.10.1", "reqwest", "ring", @@ -4448,6 +4596,40 @@ version = "11.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e" +[[package]] +name = "opendal" +version = "0.59.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f950151f9587a51a7bed70a15fa0cff464eae96e41ae7499f97067bdafdf43eb" +dependencies = [ + "opendal-core", +] + +[[package]] +name = "opendal-core" +version = "0.59.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a43405d217dfdfb543f58847336d3af672897dd1939bb7dcf314b63cf364f1c9" +dependencies = [ + "anyhow", + "asyncband", + "base64 0.23.1", + "bytes", + "futures", + "http 1.4.0", + "jiff", + "log", + "md-5 0.11.0", + "percent-encoding", + "quick-xml 0.41.0", + "serde", + "serde_json", + "tokio", + "url", + "uuid", + "web-time", +] + [[package]] name = "openssl-probe" version = "0.2.1" @@ -4530,7 +4712,6 @@ dependencies = [ "num-bigint 0.5.1", "num-integer", "num-traits", - "object_store", "ring", "seq-macro", "simdutf8", @@ -4951,6 +5132,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.11" @@ -5169,7 +5360,7 @@ version = "0.5.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" dependencies = [ - "bitflags", + "bitflags 2.11.1", ] [[package]] @@ -5178,7 +5369,7 @@ version = "0.7.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f450ad9c3b1da563fb6948a8e0fb0fb9269711c9c73d9ea1de5058c79c8d643a" dependencies = [ - "bitflags", + "bitflags 2.11.1", ] [[package]] @@ -5380,7 +5571,7 @@ version = "1.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" dependencies = [ - "bitflags", + "bitflags 2.11.1", "errno", "libc", "linux-raw-sys", @@ -5449,7 +5640,7 @@ version = "18.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "53f6a737db68eb1a8ccff86b584b2fc13eca6a7bb6f78ebc7c529547e3ab9684" dependencies = [ - "bitflags", + "bitflags 2.11.1", "cfg-if", "clipboard-win", "home", @@ -5548,7 +5739,7 @@ version = "3.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b7f4bc775c73d9a02cde8bf7b2ec4c9d12743edf609006c7facc23998404cd1d" dependencies = [ - "bitflags", + "bitflags 2.11.1", "core-foundation", "core-foundation-sys", "libc", @@ -6146,7 +6337,10 @@ dependencies = [ "arrow", "chrono-tz", "datafusion-common", + "datafusion-storage", + "datafusion-storage-object-store", "env_logger", + "object_store", "rand 0.9.4", ] @@ -6378,7 +6572,9 @@ checksum = "494815d09bf52b5548659851081238f0ca39ff638363907596da739561c62c52" dependencies = [ "bytes", "futures-core", + "futures-io", "futures-sink", + "futures-util", "libc", "pin-project-lite", "tokio", @@ -6500,7 +6696,7 @@ version = "0.6.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d4e6559d53cc268e5031cd8429d05415bc4cb4aefc4aa5d6cc35fbf5b924a1f8" dependencies = [ - "bitflags", + "bitflags 2.11.1", "bytes", "futures-util", "http 1.4.0", @@ -6790,6 +6986,7 @@ checksum = "b5772d71c9be8a8a6ac2117d949c5b224c1b72241bb611d9a3012edcf8af7812" dependencies = [ "getrandom 0.4.2", "js-sys", + "serde_core", "wasm-bindgen", ] @@ -7016,7 +7213,7 @@ version = "0.244.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "47b807c72e1bac69382b3a6fb3dbe8ea4c0ed87ff5629b8685ae6b9a611028fe" dependencies = [ - "bitflags", + "bitflags 2.11.1", "hashbrown 0.15.5", "indexmap 2.14.2", "semver", @@ -7425,7 +7622,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9d66ea20e9553b30172b5e831994e35fbde2d165325bec84fc43dbf6f4eb9cb2" dependencies = [ "anyhow", - "bitflags", + "bitflags 2.11.1", "indexmap 2.14.2", "log", "serde", diff --git a/Cargo.toml b/Cargo.toml index efeb580074138..0fff64785c1a4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -55,6 +55,9 @@ members = [ "datafusion/spark", "datafusion/sql", "datafusion/sqllogictest", + "datafusion/storage", + "datafusion/storage-object-store", + "datafusion/storage-opendal", "datafusion/substrait", "datafusion/wasmtest", "datafusion-cli", @@ -84,6 +87,10 @@ rust-version = "1.94.0" version = "55.0.0" [workspace.dependencies] +datafusion-storage-object-store = { path = "datafusion/storage-object-store", version = "55.0.0" } +datafusion-storage-opendal = { path = "datafusion/storage-opendal", version = "55.0.0" } + +datafusion-storage = { path = "datafusion/storage", version = "55.0.0", default-features = false } # We turn off default-features for some dependencies here so the workspaces which inherit them can # selectively turn them on if needed, since we can override default-features = true (from false) # for the inherited dependency but cannot do the reverse (override from true to false). @@ -123,7 +130,7 @@ criterion = "0.8" ctor = "1.0.7" dashmap = "6.2.1" datafusion = { path = "datafusion/core", version = "55.0.0", default-features = false } -datafusion-catalog = { path = "datafusion/catalog", version = "55.0.0" } +datafusion-catalog = { path = "datafusion/catalog", version = "55.0.0", default-features = false } datafusion-catalog-listing = { path = "datafusion/catalog-listing", version = "55.0.0" } datafusion-common = { path = "datafusion/common", version = "55.0.0", default-features = false } datafusion-common-runtime = { path = "datafusion/common-runtime", version = "55.0.0" } @@ -138,11 +145,19 @@ datafusion-execution = { path = "datafusion/execution", version = "55.0.0", defa datafusion-expr = { path = "datafusion/expr", version = "55.0.0", default-features = false } datafusion-expr-common = { path = "datafusion/expr-common", version = "55.0.0" } datafusion-ffi = { path = "datafusion/ffi", version = "55.0.0" } -datafusion-functions = { path = "datafusion/functions", version = "55.0.0" } -datafusion-functions-aggregate = { path = "datafusion/functions-aggregate", version = "55.0.0" } +# Select the existing expression defaults without implicitly enabling storage. +datafusion-functions = { path = "datafusion/functions", version = "55.0.0", default-features = false, features = [ + "datetime_expressions", + "encoding_expressions", + "math_expressions", + "regex_expressions", + "string_expressions", + "unicode_expressions", +] } +datafusion-functions-aggregate = { path = "datafusion/functions-aggregate", version = "55.0.0", default-features = false } datafusion-functions-aggregate-common = { path = "datafusion/functions-aggregate-common", version = "55.0.0" } datafusion-functions-nested = { path = "datafusion/functions-nested", version = "55.0.0", default-features = false } -datafusion-functions-table = { path = "datafusion/functions-table", version = "55.0.0" } +datafusion-functions-table = { path = "datafusion/functions-table", version = "55.0.0", default-features = false } datafusion-functions-window = { path = "datafusion/functions-window", version = "55.0.0" } datafusion-functions-window-common = { path = "datafusion/functions-window-common", version = "55.0.0" } datafusion-macros = { path = "datafusion/macros", version = "55.0.0" } @@ -150,15 +165,18 @@ datafusion-optimizer = { path = "datafusion/optimizer", version = "55.0.0", defa datafusion-physical-expr = { path = "datafusion/physical-expr", version = "55.0.0", default-features = false } datafusion-physical-expr-adapter = { path = "datafusion/physical-expr-adapter", version = "55.0.0", default-features = false } datafusion-physical-expr-common = { path = "datafusion/physical-expr-common", version = "55.0.0", default-features = false } -datafusion-physical-optimizer = { path = "datafusion/physical-optimizer", version = "55.0.0" } -datafusion-physical-plan = { path = "datafusion/physical-plan", version = "55.0.0" } +datafusion-physical-optimizer = { path = "datafusion/physical-optimizer", version = "55.0.0", default-features = false } +datafusion-physical-plan = { path = "datafusion/physical-plan", version = "55.0.0", default-features = false } datafusion-proto = { path = "datafusion/proto", version = "55.0.0", default-features = false } datafusion-proto-common = { path = "datafusion/proto-common", version = "55.0.0" } datafusion-proto-models = { path = "datafusion/proto-models", version = "55.0.0" } -datafusion-pruning = { path = "datafusion/pruning", version = "55.0.0" } -datafusion-session = { path = "datafusion/session", version = "55.0.0" } +datafusion-pruning = { path = "datafusion/pruning", version = "55.0.0", default-features = false } +datafusion-session = { path = "datafusion/session", version = "55.0.0", default-features = false } datafusion-spark = { path = "datafusion/spark", version = "55.0.0" } -datafusion-sql = { path = "datafusion/sql", version = "55.0.0" } +datafusion-sql = { path = "datafusion/sql", version = "55.0.0", default-features = false, features = [ + "unicode_expressions", + "unparser", +] } datafusion-substrait = { path = "datafusion/substrait", version = "55.0.0" } doc-comment = "0.3" @@ -178,11 +196,11 @@ log = "^0.4" memchr = "2.8.1" num-traits = { version = "0.2" } object_store = { version = "0.13.2", default-features = false } +opendal = { version = "0.59.1", default-features = false } parking_lot = "0.12" parquet = { version = "59.2.0", default-features = false, features = [ "arrow", "async", - "object_store", ] } pbjson = { version = "0.9.0" } pbjson-types = "0.9" diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml index cd222c3a5e5b0..ad078aead7074 100644 --- a/benchmarks/Cargo.toml +++ b/benchmarks/Cargo.toml @@ -48,6 +48,7 @@ criterion = { workspace = true, features = ["async_tokio", "html_reports"] } datafusion = { workspace = true, default-features = true } datafusion-common = { workspace = true, default-features = true } datafusion-common-runtime = { workspace = true } +datafusion-storage-object-store = { workspace = true } env_logger = { workspace = true } futures = { workspace = true } libmimalloc-sys = { version = "0.1", optional = true } diff --git a/benchmarks/src/cancellation.rs b/benchmarks/src/cancellation.rs index 5f7fdcc43d99d..4b8664cad4b83 100644 --- a/benchmarks/src/cancellation.rs +++ b/benchmarks/src/cancellation.rs @@ -30,10 +30,10 @@ use datafusion::datasource::file_format::FileFormat; use datafusion::datasource::file_format::parquet::ParquetFormat; use datafusion::datasource::listing::{ListingOptions, ListingTableUrl}; use datafusion::execution::TaskContext; -use datafusion::execution::object_store::ObjectStoreUrl; use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; use datafusion::prelude::*; +use datafusion::storage::StorageUrl; use datafusion_common::instant::Instant; use futures::TryStreamExt; use object_store::ObjectStore; @@ -166,8 +166,14 @@ async fn datafusion(store: Arc) -> Result<()> { .with_target_partitions(4) .set_bool("datafusion.execution.parquet.pushdown_filters", true); let ctx = SessionContext::new_with_config(config); - let object_store_url = ObjectStoreUrl::parse("test:///").unwrap(); - ctx.register_object_store(object_store_url.as_ref(), Arc::clone(&store)); + let object_store_url = StorageUrl::parse("test:///").unwrap(); + ctx.register_storage( + object_store_url.as_ref(), + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + store.clone(), + )), + ) + .unwrap(); let file_format = ParquetFormat::default().with_enable_pruning(true); let listing_options = ListingOptions::new(Arc::new(file_format)) diff --git a/benchmarks/src/util/options.rs b/benchmarks/src/util/options.rs index 4a1c14674a1d0..4607ed206a98e 100644 --- a/benchmarks/src/util/options.rs +++ b/benchmarks/src/util/options.rs @@ -18,6 +18,7 @@ use std::{num::NonZeroUsize, sync::Arc}; use clap::Args; +use datafusion::storage::StorageUrl; use datafusion::{ execution::{ disk_manager::DiskManagerBuilder, @@ -25,7 +26,6 @@ use datafusion::{ FairSpillPool, GreedyMemoryPool, MemoryPool, PeakRecordingPool, TrackConsumersPool, }, - object_store::ObjectStoreUrl, runtime_env::{RuntimeEnv, RuntimeEnvBuilder}, }, prelude::SessionConfig, @@ -145,8 +145,14 @@ impl CommonOpt { if self.simulate_latency { let store: Arc = Arc::new(LatencyObjectStore::new(LocalFileSystem::new())); - let url = ObjectStoreUrl::parse("file:///")?; - rt.register_object_store(url.as_ref(), store); + let url = StorageUrl::parse("file:///")?; + rt.register_storage( + url.as_ref(), + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + store, + )), + ) + .unwrap(); println!( "Simulating S3-like object store latency (get: 25-200ms, list: 40-400ms)" ); diff --git a/datafusion-cli/Cargo.toml b/datafusion-cli/Cargo.toml index 50b4002541f89..581cb8ffb6591 100644 --- a/datafusion-cli/Cargo.toml +++ b/datafusion-cli/Cargo.toml @@ -35,6 +35,8 @@ default = [] backtrace = ["datafusion/backtrace"] [dependencies] +object_store = { workspace = true, features = ["aws", "gcp", "http"] } +datafusion-storage-object-store = { workspace = true } arrow = { workspace = true } async-trait = { workspace = true } aws-config = "1.8.18" @@ -48,6 +50,7 @@ datafusion = { workspace = true, features = [ "datetime_expressions", "encoding_expressions", "nested_expressions", + "object_store", "parquet", "parquet_encryption", "recursive_protection", @@ -62,7 +65,6 @@ env_logger = { workspace = true } futures = { workspace = true } log = { workspace = true } mimalloc = { version = "0.1", default-features = false } -object_store = { workspace = true, features = ["aws", "gcp", "http"] } parking_lot = { workspace = true } parquet = { workspace = true, default-features = false } regex = { workspace = true } diff --git a/datafusion-cli/examples/cli-session-context.rs b/datafusion-cli/examples/cli-session-context.rs index 6095072163870..534a51f1332cc 100644 --- a/datafusion-cli/examples/cli-session-context.rs +++ b/datafusion-cli/examples/cli-session-context.rs @@ -20,6 +20,7 @@ use std::sync::Arc; +use datafusion::storage::{Storage, StorageBinding}; use datafusion::{ dataframe::DataFrame, error::DataFusionError, @@ -29,10 +30,8 @@ use datafusion::{ }; use datafusion_cli::{ cli_context::CliSessionContext, exec::exec_from_repl, - object_storage::instrumented::InstrumentedObjectStoreRegistry, - print_options::PrintOptions, + object_storage::instrumented::ObjectStoreProfiler, print_options::PrintOptions, }; -use object_store::ObjectStore; /// This is a toy example of a custom session context that unions the input plan with itself. struct MyUnionerContext { @@ -57,12 +56,12 @@ impl CliSessionContext for MyUnionerContext { self.ctx.state() } - fn register_object_store( + fn register_storage( &self, url: &url::Url, - object_store: Arc, - ) -> Option> { - self.ctx.register_object_store(url, object_store) + storage: Arc, + ) -> Result>, DataFusionError> { + self.ctx.register_storage(url, storage) } fn register_table_options_extension_from_scheme(&self, _scheme: &str) { @@ -91,7 +90,7 @@ pub async fn main() { quiet: false, maxrows: datafusion_cli::print_options::MaxRows::Unlimited, color: true, - instrumented_registry: Arc::new(InstrumentedObjectStoreRegistry::new()), + object_store_profiler: Arc::new(ObjectStoreProfiler::new()), }; exec_from_repl(&my_ctx, &mut print_options).await.unwrap(); diff --git a/datafusion-cli/src/catalog.rs b/datafusion-cli/src/catalog.rs index ca24da7873bc1..e17dba99f94d0 100644 --- a/datafusion-cli/src/catalog.rs +++ b/datafusion-cli/src/catalog.rs @@ -17,7 +17,7 @@ use std::sync::{Arc, Weak}; -use crate::object_storage::{AwsOptions, GcpOptions, get_object_store}; +use crate::object_storage::{AwsOptions, GcpOptions, get_storage}; use datafusion::catalog::{CatalogProvider, CatalogProviderList, SchemaProvider}; @@ -163,7 +163,7 @@ impl SchemaProvider for DynamicObjectStoreSchemaProvider { // will return `Ok` which means we don't need to register it again. However, // if `get_store` returns an `Err` then it means the corresponding store is // not registered yet and we need to register it - match state.runtime_env().object_store_registry.get_store(url) { + match state.runtime_env().storage_registry.get(url) { Ok(_) => { /*Nothing to do here, store for this URL is already registered*/ } Err(_) => { // Register the store for this URL. Here we don't have access @@ -182,7 +182,7 @@ impl SchemaProvider for DynamicObjectStoreSchemaProvider { _ => {} } state = builder.build(); - let store = get_object_store( + let store = get_storage( &state, table_url.scheme(), url, @@ -190,7 +190,7 @@ impl SchemaProvider for DynamicObjectStoreSchemaProvider { false, ) .await?; - state.runtime_env().register_object_store(url, store); + state.runtime_env().register_storage(url, store)?; } } self.inner.table(name).await @@ -260,9 +260,9 @@ mod tests { // It should still create an object store for the location in the SessionState let store = ctx .runtime_env() - .object_store(ListingTableUrl::parse(location)?)?; + .storage(ListingTableUrl::parse(location)?)?; - assert_eq!(format!("{store}"), "HttpStore"); + assert_eq!(store.url().as_str(), format!("http://{domain}/")); // The store must be configured for this domain let expected_domain = format!("Domain(\"{domain}\")"); @@ -296,8 +296,8 @@ mod tests { let store = ctx .runtime_env() - .object_store(ListingTableUrl::parse(location)?)?; - assert_eq!(format!("{store}"), format!("AmazonS3({bucket})")); + .storage(ListingTableUrl::parse(location)?)?; + assert_eq!(store.url().as_str(), format!("s3://{bucket}/")); // The store must be configured for this domain let expected_bucket = format!("bucket: \"{bucket}\""); @@ -318,8 +318,8 @@ mod tests { let store = ctx .runtime_env() - .object_store(ListingTableUrl::parse(location)?)?; - assert_eq!(format!("{store}"), format!("GoogleCloudStorage({bucket})")); + .storage(ListingTableUrl::parse(location)?)?; + assert_eq!(store.url().as_str(), format!("gs://{bucket}/")); // The store must be configured for this domain let expected_bucket = format!("bucket_name_encoded: \"{bucket}\""); diff --git a/datafusion-cli/src/cli_context.rs b/datafusion-cli/src/cli_context.rs index a6320f03fe4de..7d086b09edf38 100644 --- a/datafusion-cli/src/cli_context.rs +++ b/datafusion-cli/src/cli_context.rs @@ -17,6 +17,7 @@ use std::sync::Arc; +use datafusion::storage::{Storage, StorageBinding}; use datafusion::{ dataframe::DataFrame, error::DataFusionError, @@ -24,7 +25,6 @@ use datafusion::{ logical_expr::LogicalPlan, prelude::SessionContext, }; -use object_store::ObjectStore; use crate::object_storage::{AwsOptions, GcpOptions}; @@ -38,11 +38,11 @@ pub trait CliSessionContext { fn session_state(&self) -> SessionState; /// Register an object store with the session context. - fn register_object_store( + fn register_storage( &self, url: &url::Url, - object_store: Arc, - ) -> Option>; + storage: Arc, + ) -> Result>, DataFusionError>; /// Register table options extension from scheme. fn register_table_options_extension_from_scheme(&self, scheme: &str); @@ -64,12 +64,12 @@ impl CliSessionContext for SessionContext { self.state() } - fn register_object_store( + fn register_storage( &self, url: &url::Url, - object_store: Arc, - ) -> Option> { - self.register_object_store(url, object_store) + storage: Arc, + ) -> Result>, DataFusionError> { + self.register_storage(url, storage) } fn register_table_options_extension_from_scheme(&self, scheme: &str) { diff --git a/datafusion-cli/src/command.rs b/datafusion-cli/src/command.rs index e847f7fdb501b..8aab8430ce850 100644 --- a/datafusion-cli/src/command.rs +++ b/datafusion-cli/src/command.rs @@ -131,16 +131,16 @@ impl Command { exec_datafusion_err!("Failed to parse input: {mode}. Valid options are disabled, summary, trace") )?; print_options - .instrumented_registry + .object_store_profiler .set_instrument_mode(profile_mode); println!( "ObjectStore Profile mode set to {}", - print_options.instrumented_registry.instrument_mode() + print_options.object_store_profiler.instrument_mode() ); } else { println!( "ObjectStore Profile mode is {}", - print_options.instrumented_registry.instrument_mode() + print_options.object_store_profiler.instrument_mode() ); } @@ -284,7 +284,7 @@ mod tests { use crate::{ object_storage::instrumented::{ - InstrumentedObjectStoreMode, InstrumentedObjectStoreRegistry, + InstrumentedObjectStoreMode, ObjectStoreProfiler, }, print_options::MaxRows, }; @@ -300,7 +300,7 @@ mod tests { quiet: false, maxrows: MaxRows::Unlimited, color: true, - instrumented_registry: Arc::new(InstrumentedObjectStoreRegistry::new()), + object_store_profiler: Arc::new(ObjectStoreProfiler::new()), }; let mut cmd: Command = "object_store_profiling" @@ -308,7 +308,7 @@ mod tests { .expect("expected parse to succeed"); assert!(cmd.execute(&ctx, &mut print_options).await.is_ok()); assert_eq!( - print_options.instrumented_registry.instrument_mode(), + print_options.object_store_profiler.instrument_mode(), InstrumentedObjectStoreMode::default() ); @@ -317,7 +317,7 @@ mod tests { .expect("expected parse to succeed"); assert!(cmd.execute(&ctx, &mut print_options).await.is_ok()); assert_eq!( - print_options.instrumented_registry.instrument_mode(), + print_options.object_store_profiler.instrument_mode(), InstrumentedObjectStoreMode::Summary ); @@ -326,7 +326,7 @@ mod tests { .expect("expected parse to succeed"); assert!(cmd.execute(&ctx, &mut print_options).await.is_ok()); assert_eq!( - print_options.instrumented_registry.instrument_mode(), + print_options.object_store_profiler.instrument_mode(), InstrumentedObjectStoreMode::Trace ); diff --git a/datafusion-cli/src/exec.rs b/datafusion-cli/src/exec.rs index 288ce4b7351b6..cc89108ce3c73 100644 --- a/datafusion-cli/src/exec.rs +++ b/datafusion-cli/src/exec.rs @@ -23,7 +23,7 @@ use crate::print_format::PrintFormat; use crate::{ command::{Command, OutputFormat}, helper::CliHelper, - object_storage::{get_object_store, stdin::StdinUtils}, + object_storage::{get_storage, stdin::StdinUtils}, print_options::{MaxRows, PrintOptions}, }; use datafusion::common::instant::Instant; @@ -327,10 +327,7 @@ impl StatementExecutor { let df = match ctx.execute_logical_plan(plan).await { Ok(df) => Ok(df), - Err(DataFusionError::ObjectStore(err)) - if matches!(err.as_ref(), Generic { store, source: _ } if "S3".eq_ignore_ascii_case(store)) - && self.statement_for_retry.is_some() => - { + Err(err) if is_s3_error(&err) && self.statement_for_retry.is_some() => { warn!( "S3 region is incorrect, auto-detecting the correct region (this may be slow). Consider updating your region configuration." ); @@ -419,7 +416,7 @@ async fn create_plan( // object store, registered like any other scheme in `get_object_store`. for location in &mut cmd.locations { *location = StdinUtils::rewrite_location(location, format.as_ref()); - register_object_store_and_config_extensions( + register_storage_and_config_extensions( ctx, location, &cmd.options, @@ -433,7 +430,7 @@ async fn create_plan( if let LogicalPlan::Copy(copy_to) = &mut plan { let format = config_file_type_from_str(©_to.file_type.get_ext()); - register_object_store_and_config_extensions( + register_storage_and_config_extensions( ctx, ©_to.output_url, ©_to.options, @@ -472,7 +469,7 @@ async fn create_plan( /// This function can return an error if the location parsing fails, options /// alteration fails, or if the object store cannot be retrieved and registered /// successfully. -pub(crate) async fn register_object_store_and_config_extensions( +pub(crate) async fn register_storage_and_config_extensions( ctx: &dyn CliSessionContext, location: &String, options: &HashMap, @@ -498,8 +495,28 @@ pub(crate) async fn register_object_store_and_config_extensions( } table_options.alter_with_string_hash_map(options)?; + // Keep application registrations unless the statement explicitly changes + // backend options or retries SDK region detection. + let changes_backend = options + .keys() + .any(|key| key.starts_with("aws.") || key.starts_with("gcp.")); + if !resolve_region + && !changes_backend + && ctx + .session_state() + .runtime_env() + .storage_registry + .get(url) + .is_ok() + { + if scheme == StdinUtils::SCHEME { + StdinUtils::get_or_create(&ctx.session_state(), url).await?; + } + return Ok(()); + } + // Retrieve the appropriate object store based on the scheme, URL, and modified table options - let store = get_object_store( + let store = get_storage( &ctx.session_state(), scheme, url, @@ -509,14 +526,28 @@ pub(crate) async fn register_object_store_and_config_extensions( .await?; // Register the retrieved object store in the session context's runtime environment - ctx.register_object_store(url, store); + ctx.register_storage(url, store)?; Ok(()) } +/// SDK errors can be wrapped by storage and query execution layers. +fn is_s3_error(error: &DataFusionError) -> bool { + let mut source: Option<&(dyn std::error::Error + 'static)> = Some(error); + while let Some(error) = source { + if matches!(error.downcast_ref::(), Some(Generic { store, .. }) if "S3".eq_ignore_ascii_case(store)) + { + return true; + } + source = error.source(); + } + false +} + #[cfg(test)] mod tests { use super::*; + use std::sync::Arc; use datafusion::common::plan_err; @@ -524,6 +555,30 @@ mod tests { use datafusion_common::assert_contains; use url::Url; + #[tokio::test] + async fn statement_preserves_an_application_storage_registration() -> Result<()> { + let ctx = SessionContext::new(); + let url = Url::parse("s3://registered").unwrap(); + ctx.register_storage( + &url, + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + Arc::new(object_store::memory::InMemory::new()), + )), + )?; + let before = ctx.runtime_env().storage_registry.get(&url)?; + register_storage_and_config_extensions( + &ctx, + &"s3://registered/data.csv".to_string(), + &HashMap::new(), + Some(ConfigFileType::CSV), + false, + ) + .await?; + let after = ctx.runtime_env().storage_registry.get(&url)?; + assert!(Arc::ptr_eq(&before, &after)); + Ok(()) + } + async fn create_external_table_test(location: &str, sql: &str) -> Result<()> { let ctx = SessionContext::new(); let plan = ctx.state().create_logical_plan(sql).await?; @@ -531,7 +586,7 @@ mod tests { if let LogicalPlan::Ddl(DdlStatement::CreateExternalTable(cmd)) = &plan { let format = config_file_type_from_str(&cmd.file_type); for location in &cmd.locations { - register_object_store_and_config_extensions( + register_storage_and_config_extensions( &ctx, location, &cmd.options, @@ -546,7 +601,7 @@ mod tests { // Ensure the URL is supported by the object store ctx.runtime_env() - .object_store(ListingTableUrl::parse(location)?)?; + .storage(ListingTableUrl::parse(location)?)?; Ok(()) } @@ -559,7 +614,7 @@ mod tests { if let LogicalPlan::Copy(cmd) = &plan { let format = config_file_type_from_str(&cmd.file_type.get_ext()); - register_object_store_and_config_extensions( + register_storage_and_config_extensions( &ctx, &cmd.output_url, &cmd.options, @@ -573,7 +628,7 @@ mod tests { // Ensure the URL is supported by the object store ctx.runtime_env() - .object_store(ListingTableUrl::parse(location)?)?; + .storage(ListingTableUrl::parse(location)?)?; Ok(()) } @@ -628,8 +683,8 @@ mod tests { assert_eq!(copy_to.output_url, location); assert_eq!(copy_to.file_type.get_ext(), "parquet".to_string()); ctx.runtime_env() - .object_store_registry - .get_store(&Url::parse(©_to.output_url).unwrap())?; + .storage_registry + .get(&Url::parse(©_to.output_url).unwrap())?; } else { return plan_err!("LogicalPlan is not a CopyTo"); } diff --git a/datafusion-cli/src/main.rs b/datafusion-cli/src/main.rs index f82206a5bd184..87a68e2965cd9 100644 --- a/datafusion-cli/src/main.rs +++ b/datafusion-cli/src/main.rs @@ -36,7 +36,7 @@ use datafusion_cli::functions::{ ListFilesCacheFunc, MetadataCacheFunc, ParquetMetadataFunc, StatisticsCacheFunc, }; use datafusion_cli::object_storage::instrumented::{ - InstrumentedObjectStoreMode, InstrumentedObjectStoreRegistry, + InstrumentedObjectStoreMode, ObjectStoreProfiler, }; use datafusion_cli::object_storage::{StdinCarriesCommands, is_stdin_location}; use datafusion_cli::{ @@ -244,13 +244,19 @@ async fn main_inner() -> Result<()> { rt_builder = rt_builder.with_disk_manager_builder(builder); } - let instrumented_registry = Arc::new( - InstrumentedObjectStoreRegistry::new() - .with_profile_mode(args.object_store_profiling), + let object_store_profiler = Arc::new( + ObjectStoreProfiler::new().with_profile_mode(args.object_store_profiling), ); - rt_builder = rt_builder.with_object_store_registry(instrumented_registry.clone()); + let session_config = session_config.with_extension(object_store_profiler.clone()); let runtime_env = rt_builder.build_arc()?; + runtime_env.register_storage( + datafusion::storage::StorageUrl::local_filesystem().as_ref(), + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + object_store_profiler + .instrument(Arc::new(object_store::local::LocalFileSystem::new())), + )), + )?; let mut state_builder = SessionStateBuilder::new() .with_config(session_config) @@ -299,7 +305,7 @@ async fn main_inner() -> Result<()> { quiet: args.quiet, maxrows: args.maxrows, color: args.color, - instrumented_registry: Arc::clone(&instrumented_registry), + object_store_profiler: Arc::clone(&object_store_profiler), }; let repl_mode = args.repl_mode(); @@ -795,10 +801,12 @@ mod tests { let ctx = SessionContext::new_with_config_rt(SessionConfig::default(), rt); - ctx.register_object_store( + ctx.register_storage( &Url::parse("mem://test_table").unwrap(), - Arc::new(InMemory::new()), - ); + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + Arc::new(InMemory::new()), + )), + )?; ctx.register_udtf( "list_files_cache", diff --git a/datafusion-cli/src/object_storage.rs b/datafusion-cli/src/object_storage.rs index 5e6337e303f6f..6b60ff44cdf7b 100644 --- a/datafusion-cli/src/object_storage.rs +++ b/datafusion-cli/src/object_storage.rs @@ -139,7 +139,9 @@ async fn get_s3_object_store_builder_inner( .is_none() || resolve_region { - let region = resolve_bucket_region(bucket_name, &ClientOptions::new()).await?; + let region = resolve_bucket_region(bucket_name, &ClientOptions::new()) + .await + .map_err(|e| DataFusionError::External(Box::new(e)))?; builder = builder.with_region(region); } @@ -189,7 +191,7 @@ impl CredentialsFromConfig { let credentials = config .credentials_provider() .ok_or_else(|| { - DataFusionError::ObjectStore(Box::new(Generic { + DataFusionError::External(Box::new(Generic { store: "S3", source: "Failed to get S3 credentials aws_config".into(), })) @@ -219,7 +221,7 @@ impl CredentialsFromConfig { "Error getting credentials from provider: {e}{source_message}", ); - return Err(DataFusionError::ObjectStore(Box::new(Generic { + return Err(DataFusionError::External(Box::new(Generic { store: "S3", source: message.into(), }))); @@ -523,13 +525,13 @@ impl ConfigExtension for GcpOptions { const PREFIX: &'static str = "gcp"; } -pub(crate) async fn get_object_store( +pub(crate) async fn get_storage( state: &SessionState, scheme: &str, url: &Url, table_options: &TableOptions, resolve_region: bool, -) -> Result, DataFusionError> { +) -> Result, DataFusionError> { let store: Arc = match scheme { "s3" => { let Some(options) = table_options.extensions.get::() else { @@ -539,7 +541,11 @@ pub(crate) async fn get_object_store( }; let builder = get_s3_object_store_builder(url, options, resolve_region).await?; - Arc::new(builder.build()?) + Arc::new( + builder + .build() + .map_err(|e| DataFusionError::External(Box::new(e)))?, + ) } "oss" => { let Some(options) = table_options.extensions.get::() else { @@ -548,7 +554,11 @@ pub(crate) async fn get_object_store( ); }; let builder = get_oss_object_store_builder(url, options)?; - Arc::new(builder.build()?) + Arc::new( + builder + .build() + .map_err(|e| DataFusionError::External(Box::new(e)))?, + ) } "cos" => { let Some(options) = table_options.extensions.get::() else { @@ -557,7 +567,11 @@ pub(crate) async fn get_object_store( ); }; let builder = get_cos_object_store_builder(url, options)?; - Arc::new(builder.build()?) + Arc::new( + builder + .build() + .map_err(|e| DataFusionError::External(Box::new(e)))?, + ) } "gs" | "gcs" => { let Some(options) = table_options.extensions.get::() else { @@ -566,29 +580,41 @@ pub(crate) async fn get_object_store( ); }; let builder = get_gcs_object_store_builder(url, options)?; - Arc::new(builder.build()?) + Arc::new( + builder + .build() + .map_err(|e| DataFusionError::External(Box::new(e)))?, + ) } "http" | "https" => Arc::new( HttpBuilder::new() .with_client_options(ClientOptions::new().with_allow_http(true)) .with_url(url.origin().ascii_serialization()) - .build()?, + .build() + .map_err(|e| DataFusionError::External(Box::new(e)))?, ), _ if scheme == stdin::StdinUtils::SCHEME => { - stdin::StdinUtils::get_or_create(state, url).await? + return stdin::StdinUtils::get_or_create(state, url).await; } _ => { - // For other types, try to get from `object_store_registry`: - state + return Ok(state .runtime_env() - .object_store_registry - .get_store(url) - .map_err(|_| { - exec_datafusion_err!("Unsupported object store scheme: {}", scheme) - })? + .storage_registry + .get(url)? + .storage() + .clone()); } }; - Ok(store) + let store = match state + .config() + .get_extension::() + { + Some(profiler) => profiler.instrument(store), + None => store, + }; + Ok(Arc::new( + datafusion_storage_object_store::ObjectStoreStorage::new(store), + )) } #[cfg(test)] diff --git a/datafusion-cli/src/object_storage/instrumented.rs b/datafusion-cli/src/object_storage/instrumented.rs index a0321cacb374b..ae708a06ef187 100644 --- a/datafusion-cli/src/object_storage/instrumented.rs +++ b/datafusion-cli/src/object_storage/instrumented.rs @@ -33,7 +33,6 @@ use chrono::Utc; use datafusion::{ common::{HashMap, instant::Instant}, error::DataFusionError, - execution::object_store::{DefaultObjectStoreRegistry, ObjectStoreRegistry}, }; use futures::stream::{BoxStream, Stream}; use futures::{StreamExt, TryStreamExt}; @@ -43,7 +42,6 @@ use object_store::{ PutResult, Result, path::Path, }; use parking_lot::{Mutex, RwLock}; -use url::Url; /// A stream wrapper that measures the time until the first response(item or end of stream) is yielded. /// @@ -754,24 +752,21 @@ impl Default for Stats { /// Provides access to [`InstrumentedObjectStore`] instances that record requests for reporting #[derive(Debug)] -pub struct InstrumentedObjectStoreRegistry { - inner: Arc, +pub struct ObjectStoreProfiler { instrument_mode: AtomicU8, stores: RwLock>>, } -impl Default for InstrumentedObjectStoreRegistry { +impl Default for ObjectStoreProfiler { fn default() -> Self { Self::new() } } -impl InstrumentedObjectStoreRegistry { - /// Returns a new [`InstrumentedObjectStoreRegistry`] that wraps the provided - /// [`ObjectStoreRegistry`] +impl ObjectStoreProfiler { + /// Returns an empty collection of profiled object stores. pub fn new() -> Self { Self { - inner: Arc::new(DefaultObjectStoreRegistry::new()), instrument_mode: AtomicU8::new(InstrumentedObjectStoreMode::default() as u8), stores: RwLock::new(Vec::new()), } @@ -783,18 +778,18 @@ impl InstrumentedObjectStoreRegistry { } /// Provides access to all of the [`InstrumentedObjectStore`]s managed by this - /// [`InstrumentedObjectStoreRegistry`] + /// [`ObjectStoreProfiler`] pub fn stores(&self) -> Vec> { self.stores.read().clone() } /// Returns the current [`InstrumentedObjectStoreMode`] for this - /// [`InstrumentedObjectStoreRegistry`] + /// [`ObjectStoreProfiler`] pub fn instrument_mode(&self) -> InstrumentedObjectStoreMode { self.instrument_mode.load(Ordering::Relaxed).into() } - /// Sets the [`InstrumentedObjectStoreMode`] for this [`InstrumentedObjectStoreRegistry`] + /// Sets the [`InstrumentedObjectStoreMode`] for this [`ObjectStoreProfiler`] pub fn set_instrument_mode(&self, mode: InstrumentedObjectStoreMode) { self.instrument_mode.store(mode as u8, Ordering::Relaxed); for s in self.stores.read().iter() { @@ -803,28 +798,14 @@ impl InstrumentedObjectStoreRegistry { } } -impl ObjectStoreRegistry for InstrumentedObjectStoreRegistry { - fn register_store( - &self, - url: &Url, - store: Arc, - ) -> Option> { +impl ObjectStoreProfiler { + /// Wrap a backend and retain its request statistics for CLI reporting. + pub fn instrument(&self, store: Arc) -> Arc { let mode = self.instrument_mode.load(Ordering::Relaxed); let instrumented = Arc::new(InstrumentedObjectStore::new(store, AtomicU8::new(mode))); self.stores.write().push(Arc::clone(&instrumented)); - self.inner.register_store(url, instrumented) - } - - fn deregister_store( - &self, - url: &Url, - ) -> datafusion::common::Result> { - self.inner.deregister_store(url) - } - - fn get_store(&self, url: &Url) -> datafusion::common::Result> { - self.inner.get_store(url) + instrumented } } @@ -868,8 +849,8 @@ mod tests { } #[test] - fn instrumented_registry() { - let mut reg = InstrumentedObjectStoreRegistry::new(); + fn object_store_profiler() { + let mut reg = ObjectStoreProfiler::new(); assert!(reg.stores().is_empty()); assert_eq!( reg.instrument_mode(), @@ -880,12 +861,7 @@ mod tests { assert_eq!(reg.instrument_mode(), InstrumentedObjectStoreMode::Trace); let store = object_store::memory::InMemory::new(); - let url = "mem://test".parse().unwrap(); - let registered = reg.register_store(&url, Arc::new(store)); - assert!(registered.is_none()); - - let fetched = reg.get_store(&url); - assert!(fetched.is_ok()); + let _instrumented = reg.instrument(Arc::new(store)); assert_eq!(reg.stores().len(), 1); } diff --git a/datafusion-cli/src/object_storage/stdin.rs b/datafusion-cli/src/object_storage/stdin.rs index 9e63068f4f32f..e355dec6b5fb6 100644 --- a/datafusion-cli/src/object_storage/stdin.rs +++ b/datafusion-cli/src/object_storage/stdin.rs @@ -118,16 +118,20 @@ impl StdinUtils { pub(crate) async fn get_or_create( state: &SessionState, url: &Url, - ) -> Result> { - let Ok(existing) = state.runtime_env().object_store_registry.get_store(url) - else { - return Self::object_store(state, url).await; + ) -> Result> { + let Ok(existing) = state.runtime_env().storage_registry.get(url) else { + return Ok(Arc::new( + datafusion_storage_object_store::ObjectStoreStorage::new( + Self::object_store(state, url).await?, + ), + )); }; - let path = ObjectStorePath::from_url_path(url.path())?; - if existing.head(&path).await.is_err() { + let path = datafusion::storage::path::Path::from_url_path(url.path())?; + let context = datafusion::storage::FileAccessContext::new("stdin"); + if existing.stat(&path, &context).await.is_err() { let buffered = existing - .list(None) + .list(&datafusion::storage::path::Path::default(), context) .try_next() .await .ok() @@ -140,7 +144,7 @@ impl StdinUtils { STORED AS format" )); } - Ok(existing) + Ok(existing.storage().clone()) } /// Builds the object store backing the `stdin://` scheme by reading all of @@ -191,8 +195,14 @@ impl StdinUtils { ) -> Result> { let store = InMemory::new(); store - .put(&ObjectStorePath::from_url_path(url.path())?, data.into()) - .await?; + .put( + &ObjectStorePath::from_url_path(url.path()).map_err(|e| { + datafusion::common::DataFusionError::External(Box::new(e)) + })?, + data.into(), + ) + .await + .map_err(|e| datafusion::common::DataFusionError::External(Box::new(e)))?; Ok(Arc::new(store)) } } @@ -252,7 +262,12 @@ mod tests { let store = StdinUtils::in_memory_object_store(&url, data).await?; let ctx = SessionContext::new(); - ctx.register_object_store(&url, store); + ctx.register_storage( + &url, + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + store, + )), + )?; ctx.sql(&format!( "CREATE EXTERNAL TABLE t STORED AS {stored_as} LOCATION '{location}' {options}" )) @@ -276,19 +291,29 @@ mod tests { // lookup through `get_or_create` and assert it hands back that exact // store rather than rebuilding it. let url = Url::parse("stdin:///stdin.csv").unwrap(); - let path = ObjectStorePath::from_url_path(url.path())?; + let path = ObjectStorePath::from_url_path(url.path()) + .map_err(|e| datafusion::common::DataFusionError::External(Box::new(e)))?; let buffered: Arc = Arc::new(InMemory::new()); - buffered.put(&path, b"a\n1\n2\n".to_vec().into()).await?; + buffered + .put(&path, b"a\n1\n2\n".to_vec().into()) + .await + .unwrap(); let ctx = SessionContext::new(); - ctx.register_object_store(&url, Arc::clone(&buffered)); + ctx.register_storage( + &url, + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + buffered.clone(), + )), + )?; let reused = StdinUtils::get_or_create(&ctx.state(), &url).await?; - assert!( - Arc::ptr_eq(&buffered, &reused), - "get_or_create must reuse the registered stdin store, not rebuild it" - ); - let bytes = reused.get(&path).await?.bytes().await?; + let original = ctx.runtime_env().storage_registry.get(&url)?; + assert!(Arc::ptr_eq(original.storage(), &reused)); + let path = datafusion::storage::path::Path::from(path.as_ref()); + let context = datafusion::storage::FileAccessContext::default(); + let reader = reused.open(&path, context).await?; + let bytes = datafusion::storage::collect_bytes(reader.stream(None)).await?; assert_eq!(bytes.as_ref(), b"a\n1\n2\n"); Ok(()) } @@ -304,7 +329,12 @@ mod tests { StdinUtils::in_memory_object_store(&csv_url, b"a\n1\n".to_vec()).await?; let ctx = SessionContext::new(); - ctx.register_object_store(&csv_url, store); + ctx.register_storage( + &csv_url, + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + store, + )), + )?; let json_url = Url::parse("stdin:///stdin.json").unwrap(); let err = StdinUtils::get_or_create(&ctx.state(), &json_url) diff --git a/datafusion-cli/src/print_options.rs b/datafusion-cli/src/print_options.rs index 5367b5e65c0fb..110254cfb7384 100644 --- a/datafusion-cli/src/print_options.rs +++ b/datafusion-cli/src/print_options.rs @@ -22,7 +22,7 @@ use std::str::FromStr; use std::sync::Arc; use crate::object_storage::instrumented::{ - InstrumentedObjectStoreMode, InstrumentedObjectStoreRegistry, RequestSummaries, + InstrumentedObjectStoreMode, ObjectStoreProfiler, RequestSummaries, }; use crate::print_format::PrintFormat; @@ -81,7 +81,7 @@ pub struct PrintOptions { pub quiet: bool, pub maxrows: MaxRows, pub color: bool, - pub instrumented_registry: Arc, + pub object_store_profiler: Arc, } // Returns the query execution details formatted @@ -190,10 +190,10 @@ impl PrintOptions { if !self.quiet { writeln!(writer, "{formatted_exec_details}")?; - let instrument_mode = self.instrumented_registry.instrument_mode(); + let instrument_mode = self.object_store_profiler.instrument_mode(); if instrument_mode != InstrumentedObjectStoreMode::Disabled { writeln!(writer, "{OBJECT_STORE_PROFILING_HEADER}")?; - for store in self.instrumented_registry.stores() { + for store in self.object_store_profiler.stores() { let requests = store.take_requests(); if !requests.is_empty() { @@ -226,13 +226,13 @@ mod tests { #[test] fn write_output() -> Result<()> { - let instrumented_registry = Arc::new(InstrumentedObjectStoreRegistry::new()); + let object_store_profiler = Arc::new(ObjectStoreProfiler::new()); let mut print_options = PrintOptions { format: PrintFormat::Automatic, quiet: true, maxrows: MaxRows::Unlimited, color: true, - instrumented_registry: Arc::clone(&instrumented_registry), + object_store_profiler: Arc::clone(&object_store_profiler), }; let mut print_output: Vec = Vec::new(); @@ -251,7 +251,7 @@ mod tests { // clear the previous data from the output so it doesn't pollute the next test print_output.clear(); print_options - .instrumented_registry + .object_store_profiler .set_instrument_mode(InstrumentedObjectStoreMode::Trace); print_options.write_output(&mut print_output, &exec_out)?; let out_str: String = print_output diff --git a/datafusion-examples/Cargo.toml b/datafusion-examples/Cargo.toml index 6d6d917ac46ec..92b48f30659e4 100644 --- a/datafusion-examples/Cargo.toml +++ b/datafusion-examples/Cargo.toml @@ -40,6 +40,7 @@ arrow = { workspace = true } arrow-schema = { workspace = true } datafusion = { workspace = true, default-features = true, features = ["parquet_encryption"] } datafusion-common = { workspace = true } +datafusion-storage-object-store = { workspace = true } nom = "8.0.0" tempfile = { workspace = true } tokio = { workspace = true, features = ["rt-multi-thread", "parking_lot", "fs"] } diff --git a/datafusion-examples/examples/custom_data_source/adapter_serialization.rs b/datafusion-examples/examples/custom_data_source/adapter_serialization.rs index d8dba6592966c..c4d62807356b5 100644 --- a/datafusion-examples/examples/custom_data_source/adapter_serialization.rs +++ b/datafusion-examples/examples/custom_data_source/adapter_serialization.rs @@ -47,11 +47,11 @@ use datafusion::datasource::physical_plan::{FileScanConfig, FileScanConfigBuilde use datafusion::datasource::source::DataSourceExec; use datafusion::execution::TaskContext; use datafusion::execution::context::SessionContext; -use datafusion::execution::object_store::ObjectStoreUrl; use datafusion::parquet::arrow::ArrowWriter; use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::SessionConfig; +use datafusion::storage::StorageUrl; use datafusion_physical_expr_adapter::{ DefaultPhysicalExprAdapterFactory, PhysicalExprAdapter, PhysicalExprAdapterFactory, }; @@ -99,10 +99,12 @@ pub async fn adapter_serialization() -> Result<()> { let mut cfg = SessionConfig::new(); cfg.options_mut().execution.parquet.pushdown_filters = true; let ctx = SessionContext::new_with_config(cfg); - ctx.runtime_env().register_object_store( - ObjectStoreUrl::parse("memory://")?.as_ref(), - Arc::clone(&store), - ); + ctx.runtime_env().register_storage( + StorageUrl::parse("memory://")?.as_ref(), + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + store.clone(), + )), + )?; // Create a table with our custom MetadataAdapterFactory let adapter_factory = Arc::new(MetadataAdapterFactory::new("v1")); @@ -450,7 +452,10 @@ async fn write_parquet( writer.close()?; let payload = PutPayload::from_bytes(buf.into()); - store.put(path, payload).await?; + store + .put(path, payload) + .await + .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e)))?; Ok(()) } diff --git a/datafusion-examples/examples/custom_data_source/csv_json_opener.rs b/datafusion-examples/examples/custom_data_source/csv_json_opener.rs index 51c0e2167053e..c747cb913a645 100644 --- a/datafusion-examples/examples/custom_data_source/csv_json_opener.rs +++ b/datafusion-examples/examples/custom_data_source/csv_json_opener.rs @@ -21,12 +21,14 @@ use std::sync::Arc; use arrow::datatypes::{DataType, Field, Schema}; use datafusion::common::config::CsvOptions; +use datafusion::storage::{ + FileAccessContext, ObjectStoreStorage, StorageBinding, StorageUrl, +}; use datafusion::{ assert_batches_eq, datasource::{ file_format::file_compression_type::FileCompressionType, listing::PartitionedFile, - object_store::ObjectStoreUrl, physical_plan::{ CsvSource, FileSource, FileStreamBuilder, JsonOpener, JsonSource, }, @@ -51,7 +53,10 @@ pub async fn csv_json_opener() -> Result<()> { } async fn csv_opener() -> Result<()> { - let object_store = Arc::new(LocalFileSystem::new()); + let storage = Arc::new(StorageBinding::new( + StorageUrl::local_filesystem(), + Arc::new(ObjectStoreStorage::new(Arc::new(LocalFileSystem::new()))), + )); let dataset = ExampleDataset::Cars; let csv_path = dataset.path(); @@ -69,17 +74,21 @@ async fn csv_opener() -> Result<()> { .with_comment(Some(b'#')) .with_batch_size(8192); - let scan_config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source) - .with_projection_indices(Some(vec![0, 1]))? - .with_limit(Some(5)) - .with_file(PartitionedFile::new(csv_path.display().to_string(), 10)) - .build(); - - let opener = - scan_config - .file_source() - .create_file_opener(object_store, &scan_config, 0)?; + let scan_config = FileScanConfigBuilder::new(StorageUrl::local_filesystem(), source) + .with_projection_indices(Some(vec![0, 1]))? + .with_limit(Some(5)) + .with_file(PartitionedFile::new( + csv_path.display().to_string(), + std::fs::metadata(&csv_path)?.len(), + )) + .build(); + + let opener = scan_config.file_source().create_file_opener( + storage, + &scan_config, + 0, + FileAccessContext::default(), + )?; let mut result = vec![]; let metrics = ExecutionPlanMetricsSet::new(); @@ -117,7 +126,11 @@ async fn json_opener() -> Result<()> { {"num":4,"str":"foo"}"#, ); - object_store.put(&path, data.into()).await?; + let file_size = data.len() as u64; + object_store + .put(&path, data.into()) + .await + .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e)))?; let schema = Arc::new(Schema::new(vec![ Field::new("num", DataType::Int64, false), @@ -130,17 +143,20 @@ async fn json_opener() -> Result<()> { 8192, projected, FileCompressionType::UNCOMPRESSED, - Arc::new(object_store), + Arc::new(StorageBinding::new( + StorageUrl::local_filesystem(), + Arc::new(ObjectStoreStorage::new(Arc::new(object_store))), + )), true, ); let scan_config = FileScanConfigBuilder::new( - ObjectStoreUrl::local_filesystem(), + StorageUrl::local_filesystem(), Arc::new(JsonSource::new(schema)), ) .with_projection_indices(Some(vec![1, 0]))? .with_limit(Some(5)) - .with_file(PartitionedFile::new(path.to_string(), 10)) + .with_file(PartitionedFile::new(path.to_string(), file_size)) .build(); let metrics = ExecutionPlanMetricsSet::new(); diff --git a/datafusion-examples/examples/custom_data_source/custom_file_casts.rs b/datafusion-examples/examples/custom_data_source/custom_file_casts.rs index 202c0a71257e9..89665b88cde1d 100644 --- a/datafusion-examples/examples/custom_data_source/custom_file_casts.rs +++ b/datafusion-examples/examples/custom_data_source/custom_file_casts.rs @@ -30,11 +30,11 @@ use datafusion::datasource::listing::{ ListingTable, ListingTableConfig, ListingTableConfigExt, ListingTableUrl, }; use datafusion::execution::context::SessionContext; -use datafusion::execution::object_store::ObjectStoreUrl; use datafusion::parquet::arrow::ArrowWriter; use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_expr::expressions::CastExpr; use datafusion::prelude::SessionConfig; +use datafusion::storage::StorageUrl; use datafusion_physical_expr_adapter::{ DefaultPhysicalExprAdapterFactory, PhysicalExprAdapter, PhysicalExprAdapterFactory, }; @@ -68,8 +68,12 @@ pub async fn custom_file_casts() -> Result<()> { // Turn on filter pushdown so that the PhysicalExprAdapter is used cfg.options_mut().execution.parquet.pushdown_filters = true; let ctx = SessionContext::new_with_config(cfg); - ctx.runtime_env() - .register_object_store(ObjectStoreUrl::parse("memory://")?.as_ref(), store); + ctx.runtime_env().register_storage( + StorageUrl::parse("memory://")?.as_ref(), + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + store, + )), + )?; // Register our good and bad files via ListingTable let listing_table_config = @@ -136,7 +140,10 @@ async fn write_data( writer.close()?; let payload = PutPayload::from_bytes(buf.into()); - store.put(path, payload).await?; + store + .put(path, payload) + .await + .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e)))?; Ok(()) } diff --git a/datafusion-examples/examples/custom_data_source/custom_file_format.rs b/datafusion-examples/examples/custom_data_source/custom_file_format.rs index 0cfbe11877e4d..2a9260115126e 100644 --- a/datafusion-examples/examples/custom_data_source/custom_file_format.rs +++ b/datafusion-examples/examples/custom_data_source/custom_file_format.rs @@ -42,7 +42,7 @@ use datafusion::{ prelude::SessionContext, }; -use object_store::{ObjectMeta, ObjectStore}; +use datafusion::storage::{FileInfo, StorageBinding}; use tempfile::tempdir; /// Example of a custom file format that reads and writes TSV files. @@ -123,8 +123,8 @@ impl FileFormat for TSVFileFormat { async fn infer_schema( &self, state: &dyn Session, - store: &Arc, - objects: &[ObjectMeta], + store: &Arc, + objects: &[FileInfo], ) -> Result { self.csv_file_format .infer_schema(state, store, objects) @@ -134,9 +134,9 @@ impl FileFormat for TSVFileFormat { async fn infer_stats( &self, state: &dyn Session, - store: &Arc, + store: &Arc, table_schema: SchemaRef, - object: &ObjectMeta, + object: &FileInfo, ) -> Result { self.csv_file_format .infer_stats(state, store, table_schema, object) diff --git a/datafusion-examples/examples/custom_data_source/default_column_values.rs b/datafusion-examples/examples/custom_data_source/default_column_values.rs index d2024621aad76..8c8ef8f70b75a 100644 --- a/datafusion-examples/examples/custom_data_source/default_column_values.rs +++ b/datafusion-examples/examples/custom_data_source/default_column_values.rs @@ -32,7 +32,6 @@ use datafusion::common::{Result, ScalarValue}; use datafusion::datasource::listing::PartitionedFile; use datafusion::datasource::physical_plan::{FileScanConfigBuilder, ParquetSource}; use datafusion::execution::context::SessionContext; -use datafusion::execution::object_store::ObjectStoreUrl; use datafusion::logical_expr::utils::conjunction; use datafusion::logical_expr::{Expr, TableProviderFilterPushDown, TableType}; use datafusion::parquet::arrow::ArrowWriter; @@ -40,6 +39,7 @@ use datafusion::parquet::file::properties::WriterProperties; use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::{SessionConfig, lit}; +use datafusion::storage::StorageUrl; use datafusion_physical_expr_adapter::{ DefaultPhysicalExprAdapterFactory, PhysicalExprAdapter, PhysicalExprAdapterFactory, replace_columns_with_literals, @@ -47,7 +47,7 @@ use datafusion_physical_expr_adapter::{ use futures::StreamExt; use object_store::memory::InMemory; use object_store::path::Path; -use object_store::{ObjectStore, ObjectStoreExt, PutPayload}; +use object_store::{ObjectStoreExt, PutPayload}; // Metadata key for storing default values in field metadata const DEFAULT_VALUE_METADATA_KEY: &str = "example.default_value"; @@ -90,7 +90,10 @@ pub async fn default_column_values() -> Result<()> { }; let path = Path::from("example.parquet"); let payload = PutPayload::from_bytes(buf.into()); - store.put(&path, payload).await?; + store + .put(&path, payload) + .await + .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e)))?; // Create a custom table provider that handles missing columns with defaults let table_provider = Arc::new(DefaultValueTableProvider::new(logical_schema)); @@ -103,10 +106,12 @@ pub async fn default_column_values() -> Result<()> { // Register our table ctx.register_table("example_table", table_provider)?; - ctx.runtime_env().register_object_store( - ObjectStoreUrl::parse("memory://")?.as_ref(), - Arc::new(store), - ); + ctx.runtime_env().register_storage( + StorageUrl::parse("memory://")?.as_ref(), + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + Arc::new(store), + )), + )?; println!("\n=== Demonstrating default value injection in filter predicates ==="); let query = "SELECT id, name FROM example_table WHERE status = 'active' ORDER BY id"; @@ -233,11 +238,14 @@ impl TableProvider for DefaultValueTableProvider { .with_predicate(filter) .with_pushdown_filters(true); - let object_store_url = ObjectStoreUrl::parse("memory://")?; - let store = state.runtime_env().object_store(object_store_url)?; + let object_store_url = StorageUrl::parse("memory://")?; + let store = state.runtime_env().storage(object_store_url)?; let mut files = vec![]; - let mut listing = store.list(None); + let mut listing = store.list( + &datafusion::storage::path::Path::default(), + datafusion::storage::FileAccessContext::default(), + ); while let Some(file) = listing.next().await { if let Ok(file) = file { files.push(file); @@ -250,7 +258,7 @@ impl TableProvider for DefaultValueTableProvider { .collect(); let file_scan_config = FileScanConfigBuilder::new( - ObjectStoreUrl::parse("memory://")?, + StorageUrl::parse("memory://")?, Arc::new(parquet_source), ) .with_projection_indices(projection.map(|p| p.to_vec()))? diff --git a/datafusion-examples/examples/data_io/in_memory_object_store.rs b/datafusion-examples/examples/data_io/in_memory_object_store.rs index 9a308f06c5abd..3512f5975ec20 100644 --- a/datafusion-examples/examples/data_io/in_memory_object_store.rs +++ b/datafusion-examples/examples/data_io/in_memory_object_store.rs @@ -28,8 +28,8 @@ use std::sync::Arc; use arrow::datatypes::{DataType, Field, Schema}; use datafusion::assert_batches_eq; use datafusion::common::Result; -use datafusion::execution::object_store::ObjectStoreUrl; use datafusion::prelude::{CsvReadOptions, SessionContext}; +use datafusion::storage::StorageUrl; use object_store::memory::InMemory; use object_store::path::Path; use object_store::{ObjectStore, ObjectStoreExt, PutPayload}; @@ -41,9 +41,14 @@ use object_store::{ObjectStore, ObjectStoreExt, PutPayload}; pub async fn in_memory_object_store() -> Result<()> { let store: Arc = Arc::new(InMemory::new()); let ctx = SessionContext::new(); - let object_store_url = ObjectStoreUrl::parse("memory://")?; + let object_store_url = StorageUrl::parse("memory://")?; // Register a URL prefix to route reads through this object store. - ctx.register_object_store(object_store_url.as_ref(), Arc::clone(&store)); + ctx.register_storage( + object_store_url.as_ref(), + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + store.clone(), + )), + )?; let schema = Schema::new(vec![ Field::new("id", DataType::Int64, false), @@ -56,7 +61,8 @@ pub async fn in_memory_object_store() -> Result<()> { // Write bytes into the in-memory object store. store .put(&csv_path, PutPayload::from_static(csv_data)) - .await?; + .await + .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e)))?; // Read using the URL that matches the registered prefix. let csv = ctx .read_csv( diff --git a/datafusion-examples/examples/data_io/json_shredding.rs b/datafusion-examples/examples/data_io/json_shredding.rs index 79023d5581832..6008f15ea5b5a 100644 --- a/datafusion-examples/examples/data_io/json_shredding.rs +++ b/datafusion-examples/examples/data_io/json_shredding.rs @@ -31,7 +31,6 @@ use datafusion::datasource::listing::{ ListingTable, ListingTableConfig, ListingTableConfigExt, ListingTableUrl, }; use datafusion::execution::context::SessionContext; -use datafusion::execution::object_store::ObjectStoreUrl; use datafusion::logical_expr::{ ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility, }; @@ -41,6 +40,7 @@ use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_expr::{ScalarFunctionExpr, expressions}; use datafusion::prelude::SessionConfig; use datafusion::scalar::ScalarValue; +use datafusion::storage::StorageUrl; use datafusion_physical_expr_adapter::{ DefaultPhysicalExprAdapterFactory, PhysicalExprAdapter, PhysicalExprAdapterFactory, }; @@ -87,16 +87,21 @@ pub async fn json_shredding() -> Result<()> { }; let path = Path::from("example.parquet"); let payload = PutPayload::from_bytes(buf.into()); - store.put(&path, payload).await?; + store + .put(&path, payload) + .await + .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e)))?; // Set up query execution let mut cfg = SessionConfig::new(); cfg.options_mut().execution.parquet.pushdown_filters = true; let ctx = SessionContext::new_with_config(cfg); - ctx.runtime_env().register_object_store( - ObjectStoreUrl::parse("memory://")?.as_ref(), - Arc::new(store), - ); + ctx.runtime_env().register_storage( + StorageUrl::parse("memory://")?.as_ref(), + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + Arc::new(store), + )), + )?; // Create a custom table provider that rewrites struct field access let listing_table_config = diff --git a/datafusion-examples/examples/data_io/object_store_spill.rs b/datafusion-examples/examples/data_io/object_store_spill.rs index d7d5392f66953..b890fb95419f0 100644 --- a/datafusion-examples/examples/data_io/object_store_spill.rs +++ b/datafusion-examples/examples/data_io/object_store_spill.rs @@ -50,8 +50,10 @@ pub async fn object_store_spill() -> Result<()> { // remote spills. This example uses a local-file-backed ObjectStore for // simplicity. let tmp_dir = tempdir()?; - let store: Arc = - Arc::new(LocalFileSystem::new_with_prefix(tmp_dir.path())?); + let store: Arc = Arc::new( + LocalFileSystem::new_with_prefix(tmp_dir.path()) + .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e)))?, + ); // Create the custom TempFileFactory that creates spill files in the ObjectStore. let temp_file_factory = Arc::new(ObjectStoreTempFileFactory::new(store)); @@ -188,7 +190,7 @@ impl SpillFile for ObjectStoreSpillFile { async move { store.get(&location).await.map(|r| r.into_stream()) }; let stream = stream::once(result_stream) .try_flatten() - .map_err(Into::into); + .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e))); Ok(Box::pin(stream)) } @@ -252,7 +254,8 @@ impl SpillWriter for ObjectStoreSpillWriter { block_on_object_store(async move { store .put(&location, PutPayload::from_bytes(data.into())) - .await?; + .await + .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e)))?; Ok(()) })?; diff --git a/datafusion-examples/examples/data_io/parquet_advanced_index.rs b/datafusion-examples/examples/data_io/parquet_advanced_index.rs index 43174bd76cf6b..8697f8cdce816 100644 --- a/datafusion-examples/examples/data_io/parquet_advanced_index.rs +++ b/datafusion-examples/examples/data_io/parquet_advanced_index.rs @@ -34,7 +34,6 @@ use datafusion::datasource::physical_plan::parquet::ParquetAccessPlan; use datafusion::datasource::physical_plan::{ FileScanConfigBuilder, ParquetFileReaderFactory, ParquetSource, }; -use datafusion::execution::object_store::ObjectStoreUrl; use datafusion::logical_expr::utils::conjunction; use datafusion::logical_expr::{TableProviderFilterPushDown, TableType}; use datafusion::parquet::arrow::ArrowWriter; @@ -52,6 +51,7 @@ use datafusion::physical_optimizer::pruning::PruningPredicateBuilder; use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; use datafusion::prelude::*; +use datafusion::storage::StorageUrl; use arrow::array::{ArrayRef, Int32Array, RecordBatch, StringArray}; use arrow::datatypes::SchemaRef; @@ -174,7 +174,12 @@ pub async fn parquet_advanced_index() -> Result<()> { // register object store provider for urls like `file://` work let url = Url::try_from("file://").unwrap(); - ctx.register_object_store(&url, object_store); + ctx.register_storage( + &url, + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + object_store, + )), + )?; // Select data from the table without any predicates (and thus no pruning) println!("** Select data, no predicates:"); @@ -484,7 +489,7 @@ impl TableProvider for IndexTableProvider { // Prepare for scanning let schema = self.schema(); - let object_store_url = ObjectStoreUrl::parse("file://")?; + let object_store_url = StorageUrl::parse("file://")?; // Configure a factory interface to avoid re-reading the metadata for each file let reader_factory = @@ -579,7 +584,9 @@ impl ParquetFileReaderFactory for CachedParquetFileReaderFactory { .to_string(); let object_store = Arc::clone(&self.object_store); - let location = partitioned_file.object_meta.location; + let location = object_store::path::Path::from( + partitioned_file.object_meta.location.as_ref(), + ); let metadata = self .metadata diff --git a/datafusion-examples/examples/data_io/parquet_embedded_index.rs b/datafusion-examples/examples/data_io/parquet_embedded_index.rs index 9f205b8b8e306..edeb473520fd1 100644 --- a/datafusion-examples/examples/data_io/parquet_embedded_index.rs +++ b/datafusion-examples/examples/data_io/parquet_embedded_index.rs @@ -123,7 +123,6 @@ use datafusion::datasource::TableType; use datafusion::datasource::listing::PartitionedFile; use datafusion::datasource::memory::DataSourceExec; use datafusion::datasource::physical_plan::{FileScanConfigBuilder, ParquetSource}; -use datafusion::execution::object_store::ObjectStoreUrl; use datafusion::logical_expr::{Operator, TableProviderFilterPushDown}; use datafusion::parquet::arrow::ArrowWriter; use datafusion::parquet::errors::ParquetError; @@ -132,6 +131,7 @@ use datafusion::parquet::file::reader::{FileReader, SerializedFileReader}; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::*; use datafusion::scalar::ScalarValue; +use datafusion::storage::StorageUrl; use std::fs::{File, read_dir}; use std::io::{Read, Seek, SeekFrom, Write}; use std::path::{Path, PathBuf}; @@ -443,7 +443,7 @@ impl TableProvider for DistinctIndexTable { println!("Scanning only files: {files_to_scan:?}"); // Build ParquetSource to actually read the files - let url = ObjectStoreUrl::parse("file://")?; + let url = StorageUrl::parse("file://")?; let source = Arc::new( ParquetSource::new(self.schema.clone()).with_enable_page_index(true), ); diff --git a/datafusion-examples/examples/data_io/parquet_encrypted_with_kms.rs b/datafusion-examples/examples/data_io/parquet_encrypted_with_kms.rs index 8e92f465eafe9..a23f268b01921 100644 --- a/datafusion-examples/examples/data_io/parquet_encrypted_with_kms.rs +++ b/datafusion-examples/examples/data_io/parquet_encrypted_with_kms.rs @@ -33,8 +33,8 @@ use datafusion::parquet::encryption::{ decrypt::FileDecryptionProperties, encrypt::FileEncryptionProperties, }; use datafusion::prelude::SessionContext; +use datafusion::storage::path::Path; use futures::StreamExt; -use object_store::path::Path; use rand::rand_core::{OsRng, TryRngCore}; use std::collections::HashSet; use std::sync::Arc; diff --git a/datafusion-examples/examples/data_io/parquet_index.rs b/datafusion-examples/examples/data_io/parquet_index.rs index 8ca63516b2c08..47d80dc787592 100644 --- a/datafusion-examples/examples/data_io/parquet_index.rs +++ b/datafusion-examples/examples/data_io/parquet_index.rs @@ -33,7 +33,6 @@ use datafusion::datasource::TableProvider; use datafusion::datasource::listing::PartitionedFile; use datafusion::datasource::memory::DataSourceExec; use datafusion::datasource::physical_plan::{FileScanConfigBuilder, ParquetSource}; -use datafusion::execution::object_store::ObjectStoreUrl; use datafusion::logical_expr::{ TableProviderFilterPushDown, TableType, utils::conjunction, }; @@ -45,6 +44,7 @@ use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_optimizer::pruning::PruningPredicateBuilder; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::*; +use datafusion::storage::StorageUrl; use std::collections::HashSet; use std::fmt::Display; use std::fs; @@ -127,7 +127,12 @@ pub async fn parquet_index() -> Result<()> { // register object store provider for urls like `file://` work let url = Url::try_from("file://").unwrap(); let object_store = object_store::local::LocalFileSystem::new(); - ctx.register_object_store(&url, Arc::new(object_store)); + ctx.register_storage( + &url, + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + Arc::new(object_store), + )), + )?; // Select data from the table without any predicates (and thus no pruning) println!("** Select data, no predicates:"); @@ -237,7 +242,7 @@ impl TableProvider for IndexTableProvider { // will not be returned. let files = self.index.get_files(predicate.clone())?; - let object_store_url = ObjectStoreUrl::parse("file://")?; + let object_store_url = StorageUrl::parse("file://")?; let source = Arc::new(ParquetSource::new(self.schema()).with_predicate(predicate)); let mut file_scan_config_builder = diff --git a/datafusion-examples/examples/data_io/partitioned_file_schema.rs b/datafusion-examples/examples/data_io/partitioned_file_schema.rs index b423ebb6a38b0..c4536dad3c6eb 100644 --- a/datafusion-examples/examples/data_io/partitioned_file_schema.rs +++ b/datafusion-examples/examples/data_io/partitioned_file_schema.rs @@ -21,13 +21,13 @@ use arrow::array::{Int32Array, RecordBatch}; use arrow_schema::{DataType, Field, Schema, SchemaRef}; use datafusion::common::Result; use datafusion::datasource::listing::PartitionedFile; -use datafusion::datasource::object_store::ObjectStoreUrl; use datafusion::datasource::physical_plan::{FileScanConfigBuilder, ParquetSource}; use datafusion::datasource::source::DataSourceExec; use datafusion::execution::TaskContext; use datafusion::parquet::arrow::ArrowWriter; use datafusion::parquet::file::reader::Length; use datafusion::physical_plan::ExecutionPlan; +use datafusion::storage::StorageUrl; use futures::StreamExt; use std::fs::File; use std::path::Path; @@ -131,7 +131,7 @@ async fn read_file( } let config = FileScanConfigBuilder::new( - ObjectStoreUrl::local_filesystem(), + StorageUrl::local_filesystem(), Arc::new(ParquetSource::new(source_schema)), ) .with_file(partitioned_file) diff --git a/datafusion-examples/examples/data_io/query_http_csv.rs b/datafusion-examples/examples/data_io/query_http_csv.rs index 71421e6270ccb..e4b188d42a6ae 100644 --- a/datafusion-examples/examples/data_io/query_http_csv.rs +++ b/datafusion-examples/examples/data_io/query_http_csv.rs @@ -34,7 +34,12 @@ pub async fn query_http_csv() -> Result<()> { .with_url(base_url.clone()) .build() .unwrap(); - ctx.register_object_store(&base_url, Arc::new(http_store)); + ctx.register_storage( + &base_url, + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + Arc::new(http_store), + )), + )?; // register csv file with the execution context ctx.register_csv( diff --git a/datafusion-examples/examples/external_dependency/dataframe_to_s3.rs b/datafusion-examples/examples/external_dependency/dataframe_to_s3.rs index fdb8a3c9c051a..2e259de4f93e9 100644 --- a/datafusion-examples/examples/external_dependency/dataframe_to_s3.rs +++ b/datafusion-examples/examples/external_dependency/dataframe_to_s3.rs @@ -50,12 +50,18 @@ pub async fn dataframe_to_s3() -> Result<()> { .with_region(region) .with_access_key_id(env::var("AWS_ACCESS_KEY_ID").unwrap()) .with_secret_access_key(env::var("AWS_SECRET_ACCESS_KEY").unwrap()) - .build()?; + .build() + .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e)))?; let path = format!("s3://{bucket_name}"); let s3_url = Url::parse(&path).unwrap(); let arc_s3 = Arc::new(s3); - ctx.register_object_store(&s3_url, arc_s3.clone()); + ctx.register_storage( + &s3_url, + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + arc_s3.clone(), + )), + )?; let path = format!("s3://{bucket_name}/test_data/"); let file_format = ParquetFormat::default().with_enable_pruning(true); diff --git a/datafusion-examples/examples/external_dependency/query_aws_s3.rs b/datafusion-examples/examples/external_dependency/query_aws_s3.rs index 7dc2f76be4f0c..292459aab49dc 100644 --- a/datafusion-examples/examples/external_dependency/query_aws_s3.rs +++ b/datafusion-examples/examples/external_dependency/query_aws_s3.rs @@ -44,11 +44,17 @@ pub async fn query_aws_s3() -> Result<()> { .with_region(region) .with_access_key_id(env::var("AWS_ACCESS_KEY_ID").unwrap()) .with_secret_access_key(env::var("AWS_SECRET_ACCESS_KEY").unwrap()) - .build()?; + .build() + .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e)))?; let path = format!("s3://{bucket_name}"); let s3_url = Url::parse(&path).unwrap(); - ctx.register_object_store(&s3_url, Arc::new(s3)); + ctx.register_storage( + &s3_url, + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + Arc::new(s3), + )), + )?; // cannot query the parquet files from this bucket because the path contains a whitespace // and we don't support that yet diff --git a/datafusion-examples/examples/query_planning/thread_pools.rs b/datafusion-examples/examples/query_planning/thread_pools.rs index af3806e77d816..7480648c5e9f3 100644 --- a/datafusion-examples/examples/query_planning/thread_pools.rs +++ b/datafusion-examples/examples/query_planning/thread_pools.rs @@ -214,10 +214,16 @@ async fn different_runtime_advanced() -> Result<()> { // you will see an error such as: // A Tokio 1.x context was found, but IO is disabled. .with_http_connector(SpawnedReqwestConnector::new(io_handle)) - .build()?; + .build() + .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e)))?; // Tell DataFusion to process `http://` urls with this wrapped object store - ctx.register_object_store(&base_url, Arc::new(http_store)); + ctx.register_storage( + &base_url, + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + Arc::new(http_store), + )), + )?; // As above, plan and execute the query on the cpu runtime. let (tx, mut rx) = tokio::sync::mpsc::channel(2); diff --git a/datafusion-examples/examples/sql_ops/custom_sql_parser.rs b/datafusion-examples/examples/sql_ops/custom_sql_parser.rs index 308a0de62a242..f7c0e05b9d3ed 100644 --- a/datafusion-examples/examples/sql_ops/custom_sql_parser.rs +++ b/datafusion-examples/examples/sql_ops/custom_sql_parser.rs @@ -177,9 +177,17 @@ async fn handle_create_external_catalog( // Register a local object store rooted at the workspace root. // We use a specific authority 'workspace' to ensure consistent resolution. - let store = Arc::new(LocalFileSystem::new_with_prefix(workspace_root())?); + let store = Arc::new( + LocalFileSystem::new_with_prefix(workspace_root()) + .map_err(|e| DataFusionError::External(Box::new(e)))?, + ); let store_url = url::Url::parse("local://workspace").unwrap(); - ctx.register_object_store(&store_url, Arc::clone(&store) as _); + ctx.register_storage( + &store_url, + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + store.clone() as _, + )), + )?; let target_ext = format!(".{}", stmt.catalog_type.to_lowercase()); @@ -195,7 +203,7 @@ async fn handle_create_external_catalog( let mut list_stream = store.list(Some(&prefix)); while let Some(meta) = list_stream.next().await { - let meta = meta?; + let meta = meta.map_err(|e| DataFusionError::External(Box::new(e)))?; let path = &meta.location; if path.as_ref().ends_with(&target_ext) { diff --git a/datafusion-examples/examples/sql_ops/query.rs b/datafusion-examples/examples/sql_ops/query.rs index 60b47c36b9ae2..07ad9a0ee5a24 100644 --- a/datafusion-examples/examples/sql_ops/query.rs +++ b/datafusion-examples/examples/sql_ops/query.rs @@ -174,7 +174,12 @@ async fn query_parquet() -> Result<()> { let url = url::Url::parse("file://./") .map_err(|e| DataFusionError::External(Box::new(e)))?; - ctx.register_object_store(&url, local_fs); + ctx.register_storage( + &url, + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + local_fs, + )), + )?; // Register a listing table - this will use all files in the directory as data sources // for the query diff --git a/datafusion/catalog-listing/Cargo.toml b/datafusion/catalog-listing/Cargo.toml index abe58f45994be..bf3a7cd937f14 100644 --- a/datafusion/catalog-listing/Cargo.toml +++ b/datafusion/catalog-listing/Cargo.toml @@ -31,10 +31,11 @@ version.workspace = true all-features = true [dependencies] +datafusion-storage.workspace = true arrow = { workspace = true } async-trait = { workspace = true } datafusion-catalog = { workspace = true } -datafusion-common = { workspace = true, features = ["object_store"] } +datafusion-common = { workspace = true } datafusion-datasource = { workspace = true } datafusion-execution = { workspace = true } datafusion-expr = { workspace = true } @@ -45,10 +46,10 @@ datafusion-physical-plan = { workspace = true } futures = { workspace = true } itertools = { workspace = true } log = { workspace = true } -object_store = { workspace = true } percent-encoding = { workspace = true } [dev-dependencies] +object_store = { workspace = true } chrono = { workspace = true } datafusion-datasource-parquet = { workspace = true } diff --git a/datafusion/catalog-listing/src/config.rs b/datafusion/catalog-listing/src/config.rs index 2b83c8ec92b2c..270f0c3bc799b 100644 --- a/datafusion/catalog-listing/src/config.rs +++ b/datafusion/catalog-listing/src/config.rs @@ -53,6 +53,8 @@ pub enum SchemaSource { /// - **Custom handling of missing columns**: By default they are filled in with nulls, but you may e.g. want to fill them in with `0` or `""`. #[derive(Debug, Clone, Default)] pub struct ListingTableConfig { + /// Retained registration used for discovery and execution. + pub storage: Option>, /// Paths on the `ObjectStore` for creating [`crate::ListingTable`]. /// They should share the same schema and object store. pub table_paths: Vec, @@ -89,6 +91,15 @@ impl ListingTableConfig { } } + /// Bind this table to a registration instead of resolving its URL again. + pub fn with_storage( + mut self, + storage: Arc, + ) -> Self { + self.storage = Some(storage); + self + } + /// Returns the source of the schema for this configuration pub fn schema_source(&self) -> SchemaSource { self.schema_source @@ -203,12 +214,18 @@ impl ListingTableConfig { /// # Errors /// * if `self.options` is not set. See [`Self::with_listing_options`] pub async fn infer_schema( - self, + mut self, state: &dyn Session, ) -> datafusion_common::Result { + if self.storage.is_none() + && let Some(url) = self.table_paths.first() + { + self.storage = Some(state.runtime_env().storage(url)?); + } match self.options { Some(options) => { let ListingTableConfig { + storage, table_paths, file_schema, options: _, @@ -221,7 +238,13 @@ impl ListingTableConfig { None => { if let Some(url) = table_paths.first() { ( - options.infer_schema(state, url).await?, + options + .infer_schema_with_storage( + state, + url, + storage.as_ref().expect("table storage"), + ) + .await?, SchemaSource::Inferred, ) } else { @@ -231,6 +254,7 @@ impl ListingTableConfig { }; Ok(Self { + storage, table_paths, file_schema: Some(schema), options: Some(options), @@ -247,9 +271,14 @@ impl ListingTableConfig { /// # Errors /// * if `self.options` is not set. See [`Self::with_listing_options`] pub async fn infer_partitions_from_path( - self, + mut self, state: &dyn Session, ) -> datafusion_common::Result { + if self.storage.is_none() + && let Some(url) = self.table_paths.first() + { + self.storage = Some(state.runtime_env().storage(url)?); + } match self.options { Some(options) => { let Some(url) = self.table_paths.first() else { @@ -271,6 +300,7 @@ impl ListingTableConfig { .collect::>(); let options = options.with_table_partition_cols(partitions); Ok(Self { + storage: self.storage, table_paths: self.table_paths, file_schema: self.file_schema, options: Some(options), diff --git a/datafusion/catalog-listing/src/helpers.rs b/datafusion/catalog-listing/src/helpers.rs index dc090378a8513..dc59b28baab3e 100644 --- a/datafusion/catalog-listing/src/helpers.rs +++ b/datafusion/catalog-listing/src/helpers.rs @@ -43,8 +43,8 @@ use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion_common::{Column, DFSchema}; use datafusion_expr::{Expr, Volatility}; use datafusion_physical_expr::create_physical_expr; -use object_store::path::Path; -use object_store::{ObjectMeta, ObjectStore}; +use datafusion_storage::path::Path; +use datafusion_storage::{FileInfo, StorageBinding}; use percent_encoding::{AsciiSet, CONTROLS, percent_decode_str, utf8_percent_encode}; const PARTITION_VALUE_ENCODE_SET: &AsciiSet = @@ -150,30 +150,34 @@ pub struct Partition { /// or equivalently the number of partition values in `path` depth: usize, /// The files contained as direct children of this `Partition` if known - files: Option>, + files: Option>, } impl Partition { /// List the direct children of this partition updating `self.files` with /// any child files, and returning a list of child "directories" - async fn list(mut self, store: &dyn ObjectStore) -> Result<(Self, Vec)> { + async fn list(mut self, store: &StorageBinding) -> Result<(Self, Vec)> { trace!("Listing partition {}", self.path); - let prefix = Some(&self.path).filter(|p| !p.as_ref().is_empty()); - let result = store.list_with_delimiter(prefix).await?; + let result = store + .list_with_delimiter( + &self.path, + &datafusion_storage::FileAccessContext::new("listing"), + ) + .await?; self.files = Some( result - .objects + .files .into_iter() .filter(|object_meta| object_meta.size > 0) .collect(), ); - Ok((self, result.common_prefixes)) + Ok((self, result.directories)) } } /// Returns a recursive list of the partitions in `table_path` up to `max_depth` pub async fn list_partitions( - store: &dyn ObjectStore, + store: &StorageBinding, table_path: &ListingTableUrl, max_depth: usize, partition_prefix: Option, @@ -351,7 +355,7 @@ pub fn filter_partitioned_file( /// because hive-style partition values are never null and there is no valid /// value to assign for non-partitioned files. fn try_into_partitioned_file( - object_meta: ObjectMeta, + object_meta: FileInfo, partition_cols: &[(String, DataType)], table_path: &ListingTableUrl, ) -> Result> { @@ -384,7 +388,7 @@ fn try_into_partitioned_file( /// using only the partition columns. pub async fn pruned_partition_list<'a>( ctx: &'a dyn Session, - store: &'a dyn ObjectStore, + store: &'a StorageBinding, table_path: &'a ListingTableUrl, filters: &'a [Expr], file_extension: &'a str, @@ -442,7 +446,7 @@ pub async fn pruned_partition_list<'a>( } fn object_meta_to_partitioned_file( - object_meta: ObjectMeta, + object_meta: FileInfo, table_ref: Option<&TableReference>, ) -> Result> { Ok(Some(PartitionedFile { @@ -664,7 +668,7 @@ mod tests { fn test_try_into_partitioned_file_valid_partition() { let table_path = ListingTableUrl::parse("file:///bucket/mytable").unwrap(); let partition_cols = vec![("year_month".to_string(), DataType::Utf8)]; - let meta = ObjectMeta { + let meta = FileInfo { location: Path::from("bucket/mytable/year_month=2024-01/data.parquet"), last_modified: chrono::Utc::now(), size: 100, @@ -687,7 +691,7 @@ mod tests { fn test_try_into_partitioned_file_decodes_partition_value() { let table_path = ListingTableUrl::parse("file:///bucket/mytable").unwrap(); let partition_cols = vec![("category".to_string(), DataType::Utf8)]; - let meta = ObjectMeta { + let meta = FileInfo { location: Path::parse( "bucket/mytable/category=Electronics%2FComputers/data.parquet", ) @@ -716,7 +720,7 @@ mod tests { // hive partitioning was added. let table_path = ListingTableUrl::parse("file:///bucket/mytable").unwrap(); let partition_cols = vec![("year_month".to_string(), DataType::Utf8)]; - let meta = ObjectMeta { + let meta = FileInfo { location: Path::from("bucket/mytable/data.parquet"), last_modified: chrono::Utc::now(), size: 100, @@ -737,7 +741,7 @@ mod tests { // File in a directory that doesn't match the expected partition column let table_path = ListingTableUrl::parse("file:///bucket/mytable").unwrap(); let partition_cols = vec![("year_month".to_string(), DataType::Utf8)]; - let meta = ObjectMeta { + let meta = FileInfo { location: Path::from("bucket/mytable/wrong_col=2024-01/data.parquet"), last_modified: chrono::Utc::now(), size: 100, @@ -760,7 +764,7 @@ mod tests { ("year".to_string(), DataType::Utf8), ("month".to_string(), DataType::Utf8), ]; - let meta = ObjectMeta { + let meta = FileInfo { location: Path::from("bucket/mytable/year=2024/month=01/data.parquet"), last_modified: chrono::Utc::now(), size: 100, @@ -791,7 +795,7 @@ mod tests { ("year".to_string(), DataType::Utf8), ("month".to_string(), DataType::Utf8), ]; - let meta = ObjectMeta { + let meta = FileInfo { location: Path::from("bucket/mytable/year=2024/data.parquet"), last_modified: chrono::Utc::now(), size: 100, diff --git a/datafusion/catalog-listing/src/options.rs b/datafusion/catalog-listing/src/options.rs index 44337e52a1e05..ca3f36ee1b30b 100644 --- a/datafusion/catalog-listing/src/options.rs +++ b/datafusion/catalog-listing/src/options.rs @@ -277,8 +277,17 @@ impl ListingOptions { state: &dyn Session, table_path: &'a ListingTableUrl, ) -> datafusion_common::Result { - let store = state.runtime_env().object_store(table_path)?; + let store = state.runtime_env().storage(table_path)?; + self.infer_schema_with_storage(state, table_path, &store) + .await + } + pub(crate) async fn infer_schema_with_storage( + &self, + state: &dyn Session, + table_path: &ListingTableUrl, + store: &Arc, + ) -> datafusion_common::Result { let all_files: Vec<_> = table_path .list_all_files(state, store.as_ref(), &self.file_extension) .await? @@ -300,7 +309,7 @@ impl ListingOptions { .filter(|object_meta| object_meta.size > 0) .collect(); - let schema = self.format.infer_schema(state, &store, &files).await?; + let schema = self.format.infer_schema(state, store, &files).await?; Ok(schema) } @@ -369,7 +378,7 @@ impl ListingOptions { state: &dyn Session, table_path: &ListingTableUrl, ) -> datafusion_common::Result> { - let store = state.runtime_env().object_store(table_path)?; + let store = state.runtime_env().storage(table_path)?; // only use 10 files for inference // This can fail to detect inconsistent partition keys diff --git a/datafusion/catalog-listing/src/table.rs b/datafusion/catalog-listing/src/table.rs index 6c294fe077db4..1c358ab12ba4a 100644 --- a/datafusion/catalog-listing/src/table.rs +++ b/datafusion/catalog-listing/src/table.rs @@ -53,9 +53,9 @@ use datafusion_physical_expr_adapter::PhysicalExprAdapterFactory; use datafusion_physical_expr_common::sort_expr::LexOrdering; use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::empty::EmptyExec; +use datafusion_storage::StorageBinding; use futures::future::BoxFuture; use futures::{Stream, StreamExt, TryStreamExt, future, stream}; -use object_store::ObjectStore; use std::collections::{HashMap, HashSet}; use std::sync::Arc; @@ -72,7 +72,7 @@ pub struct ListFilesResult { /// Built in [`TableProvider`] that reads data from one or more files as a single table. /// -/// The files are read using an [`ObjectStore`] instance, for example from +/// The files are read using a [`StorageBinding`] instance, for example from /// local files or objects from AWS S3. /// /// # Features: @@ -179,6 +179,7 @@ pub struct ListFilesResult { /// ``` #[derive(Debug, Clone)] pub struct ListingTable { + storage: std::sync::OnceLock>, table_paths: Vec, /// `file_schema` contains only the columns physically stored in the data files themselves. /// - Represents the actual fields found in files like Parquet, CSV, etc. @@ -214,6 +215,14 @@ impl ListingTable { /// /// See documentation and example on [`ListingTable`] and [`ListingTableConfig`] pub fn try_new(config: ListingTableConfig) -> datafusion_common::Result { + if let Some(storage) = &config.storage + && config + .table_paths + .iter() + .any(|path| path.storage_url() != *storage.url()) + { + return plan_err!("All table paths must use the same storage registration"); + } // Extract schema_source before moving other parts of the config let schema_source = config.schema_source(); @@ -241,6 +250,10 @@ impl ListingTable { Arc::new(SchemaFingerprint::from_schema(&file_schema)); let table = Self { + storage: config + .storage + .map(std::sync::OnceLock::from) + .unwrap_or_default(), table_paths: config.table_paths, file_schema, table_schema, @@ -257,6 +270,31 @@ impl ListingTable { Ok(table) } + fn storage( + &self, + state: &dyn Session, + ) -> datafusion_common::Result> { + if let Some(storage) = self.storage.get() { + return Ok(Arc::clone(storage)); + } + let first = self + .table_paths + .first() + .ok_or_else(|| internal_datafusion_err!("No table path"))?; + let storage = state.runtime_env().storage(first)?; + for path in &self.table_paths { + if path.storage_url() != *storage.url() { + return plan_err!( + "All table paths must use the same storage registration" + ); + } + } + let _ = self.storage.set(storage); + Ok(Arc::clone( + self.storage.get().expect("table storage initialized"), + )) + } + /// Assign constraints pub fn with_constraints(mut self, constraints: Constraints) -> Self { self.constraints = constraints; @@ -724,7 +762,7 @@ impl ListingTable { }; let Some(object_store_url) = - self.table_paths.first().map(ListingTableUrl::object_store) + self.table_paths.first().map(ListingTableUrl::storage_url) else { return Ok(ScanResult::new(Arc::new(EmptyExec::new(Arc::new( Schema::empty(), @@ -733,6 +771,7 @@ impl ListingTable { let file_source = self.create_file_source(); let scan_config = FileScanConfigBuilder::new(object_store_url, file_source) + .with_storage(self.storage(state)?) .with_file_groups(partitioned_file_lists) .with_constraints(self.constraints.clone()) .with_statistics(statistics) @@ -805,7 +844,7 @@ impl ListingTable { } // Get the object store for the table path. - let store = state.runtime_env().object_store(table_path)?; + let store = self.storage(state)?; let file_list_stream = pruned_partition_list( state, @@ -824,6 +863,7 @@ impl ListingTable { // Invalidate cache entries for this table if they exist if let Some(lfc) = state.runtime_env().cache_manager.get_list_files_cache() { let key = TableScopedPath { + storage_id: store.id(), table: table_path.get_table_ref().clone(), path: table_path.prefix().clone(), }; @@ -832,8 +872,9 @@ impl ListingTable { // Sink related option, apart from format let config = FileSinkConfig { + storage: Some(store), original_url: String::default(), - object_store_url: self.table_paths()[0].object_store(), + object_store_url: self.table_paths()[0].storage_url(), table_paths: self.table_paths().clone(), file_group, output_schema: self.schema(), @@ -884,7 +925,7 @@ impl ListingTable { async fn collect_files_for_scan<'a>( &'a self, ctx: &'a dyn Session, - store: &'a Arc, + store: &'a Arc, listing_time_filters: &'a [Expr], file_limit: Option, ) -> datafusion_common::Result<(FileGroup, bool)> { @@ -946,8 +987,8 @@ impl ListingTable { ); } - let store = if let Some(url) = self.table_paths.first() { - ctx.runtime_env().object_store(url)? + let store = if let Some(_url) = self.table_paths.first() { + self.storage(ctx)? } else { return Ok(ListFilesResult { file_groups: vec![], @@ -1008,8 +1049,8 @@ impl ListingTable { ); } - let store = if let Some(url) = self.table_paths.first() { - ctx.runtime_env().object_store(url)? + let store = if let Some(_url) = self.table_paths.first() { + self.storage(ctx)? } else { return Ok(ListFilesResult { file_groups: vec![], @@ -1088,10 +1129,11 @@ impl ListingTable { async fn do_collect_statistics_and_ordering( &self, ctx: &dyn Session, - store: &Arc, + store: &Arc, part_file: &PartitionedFile, ) -> datafusion_common::Result<(Arc, Option)> { let path = TableScopedPath { + storage_id: store.id(), table: part_file.table_reference.clone(), path: part_file.object_meta.location.clone(), }; diff --git a/datafusion/catalog/Cargo.toml b/datafusion/catalog/Cargo.toml index 1009e9aee477b..1dc5082f61fe7 100644 --- a/datafusion/catalog/Cargo.toml +++ b/datafusion/catalog/Cargo.toml @@ -30,6 +30,9 @@ version.workspace = true [package.metadata.docs.rs] all-features = true +[features] +default = [] + [dependencies] arrow = { workspace = true } async-trait = { workspace = true } @@ -42,10 +45,10 @@ datafusion-expr = { workspace = true } datafusion-physical-expr = { workspace = true } datafusion-physical-plan = { workspace = true } datafusion-session = { workspace = true } +datafusion-storage.workspace = true futures = { workspace = true } itertools = { workspace = true } log = { workspace = true } -object_store = { workspace = true } parking_lot = { workspace = true } tokio = { workspace = true } @@ -54,3 +57,6 @@ tokio = { workspace = true } # https://github.com/rust-lang/cargo/issues/13157 [lints] workspace = true + +[dev-dependencies] +object_store = { workspace = true } diff --git a/datafusion/catalog/src/listing_schema.rs b/datafusion/catalog/src/listing_schema.rs index d38fe659aaa97..70a447b14d557 100644 --- a/datafusion/catalog/src/listing_schema.rs +++ b/datafusion/catalog/src/listing_schema.rs @@ -30,11 +30,11 @@ use datafusion_common::{ use datafusion_expr::CreateExternalTable; use async_trait::async_trait; +use datafusion_storage::StorageBinding; use futures::TryStreamExt; use itertools::Itertools; -use object_store::ObjectStore; -/// A [`SchemaProvider`] that scans an [`ObjectStore`] to automatically discover tables +/// A [`SchemaProvider`] that scans a [`StorageBinding`] to automatically discover tables /// /// A subfolder relationship is assumed, i.e. given: /// - authority = `s3://host.example.com:3000` @@ -48,13 +48,12 @@ use object_store::ObjectStore; /// - `s3://host.example.com:3000/data/tpch/customer/part-00000-xxxx.snappy.parquet` /// - `s3://host.example.com:3000/data/tpch/customer/_delta_log/` /// -/// [`ObjectStore`]: object_store::ObjectStore #[derive(Debug)] pub struct ListingSchemaProvider { authority: String, - path: object_store::path::Path, + path: datafusion_storage::path::Path, factory: Arc, - store: Arc, + store: Arc, tables: Arc>>>, format: String, } @@ -66,14 +65,14 @@ impl ListingSchemaProvider { /// `authority`: The scheme (i.e. s3://) + host (i.e. example.com:3000) /// `path`: The root path that contains subfolders which represent tables /// `factory`: The `TableProviderFactory` to use to instantiate tables for each subfolder - /// `store`: The `ObjectStore` containing the table data + /// `store`: The storage binding containing the table data /// `format`: The `FileFormat` of the tables /// `has_header`: Indicates whether the created external table has the has_header flag enabled pub fn new( authority: String, - path: object_store::path::Path, + path: datafusion_storage::path::Path, factory: Arc, - store: Arc, + store: Arc, format: String, ) -> Self { Self { @@ -88,7 +87,14 @@ impl ListingSchemaProvider { /// Reload table information from ObjectStore pub async fn refresh(&self, state: &dyn Session) -> datafusion_common::Result<()> { - let entries: Vec<_> = self.store.list(Some(&self.path)).try_collect().await?; + let entries: Vec<_> = self + .store + .list( + &self.path, + datafusion_storage::FileAccessContext::new("catalog"), + ) + .try_collect() + .await?; let base = Path::new(self.path.as_ref()); let mut tables = HashSet::new(); for file in entries.iter() { diff --git a/datafusion/common/Cargo.toml b/datafusion/common/Cargo.toml index 9ee199fe82f28..4062528d5cdef 100644 --- a/datafusion/common/Cargo.toml +++ b/datafusion/common/Cargo.toml @@ -73,6 +73,7 @@ arrow = { workspace = true } arrow-ipc = { workspace = true } arrow-schema = { workspace = true, features = ["canonical_extension_types"] } chrono = { workspace = true } +datafusion-storage = { workspace = true } foldhash = "0.2" half = { workspace = true } hashbrown = { workspace = true } @@ -82,7 +83,6 @@ itertools = { workspace = true } libc = "0.2.185" log = { workspace = true } num-traits = { workspace = true } -object_store = { workspace = true, optional = true } parquet = { workspace = true, optional = true, default-features = true } recursive = { workspace = true, optional = true } sqlparser = { workspace = true, optional = true } diff --git a/datafusion/common/src/error.rs b/datafusion/common/src/error.rs index d1fcb50f73492..7b1b6484dc263 100644 --- a/datafusion/common/src/error.rs +++ b/datafusion/common/src/error.rs @@ -74,9 +74,8 @@ pub enum DataFusionError { /// Error when reading / writing Parquet data. #[cfg(feature = "parquet")] ParquetError(Box), - /// Error when reading / writing to / from an object_store (e.g. S3 or LocalFile) - #[cfg(feature = "object_store")] - ObjectStore(Box), + /// Error accessing storage. + Storage(Box), /// Error when an I/O operation fails IoError(io::Error), /// Error when SQL is syntactically incorrect. @@ -393,20 +392,6 @@ impl From for DataFusionError { } } -#[cfg(feature = "object_store")] -impl From for DataFusionError { - fn from(e: object_store::Error) -> Self { - DataFusionError::ObjectStore(Box::new(e)) - } -} - -#[cfg(feature = "object_store")] -impl From for DataFusionError { - fn from(e: object_store::path::Error) -> Self { - DataFusionError::ObjectStore(Box::new(e.into())) - } -} - #[cfg(feature = "sql")] impl From for DataFusionError { fn from(e: ParserError) -> Self { @@ -443,8 +428,7 @@ impl Error for DataFusionError { DataFusionError::ArrowError(e, _) => Some(e.as_ref()), #[cfg(feature = "parquet")] DataFusionError::ParquetError(e) => Some(e.as_ref()), - #[cfg(feature = "object_store")] - DataFusionError::ObjectStore(e) => Some(e.as_ref()), + DataFusionError::Storage(e) => Some(e.as_ref()), DataFusionError::IoError(e) => Some(e), #[cfg(feature = "sql")] DataFusionError::SQL(e, _) => Some(e.as_ref()), @@ -572,8 +556,7 @@ impl DataFusionError { DataFusionError::ArrowError(_, _) => "Arrow error: ", #[cfg(feature = "parquet")] DataFusionError::ParquetError(_) => "Parquet error: ", - #[cfg(feature = "object_store")] - DataFusionError::ObjectStore(_) => "Object Store error: ", + DataFusionError::Storage(_) => "Storage error: ", DataFusionError::IoError(_) => "IO error: ", #[cfg(feature = "sql")] DataFusionError::SQL(_, _) => "SQL error: ", @@ -635,8 +618,7 @@ impl DataFusionError { DataFusionError::ExecutionJoin(ref desc) => Cow::Owned(desc.to_string()), DataFusionError::ResourcesExhausted(ref desc) => Cow::Owned(desc.to_string()), DataFusionError::External(ref desc) => Cow::Owned(desc.to_string()), - #[cfg(feature = "object_store")] - DataFusionError::ObjectStore(ref desc) => Cow::Owned(desc.to_string()), + DataFusionError::Storage(ref desc) => Cow::Owned(desc.to_string()), DataFusionError::Context(ref desc, ref err) => { Cow::Owned(format!("{desc}\ncaused by\n{}", *err)) } @@ -1193,6 +1175,17 @@ pub fn add_possible_columns_to_diag( } } +impl From for DataFusionError { + fn from(error: datafusion_storage::Error) -> Self { + Self::Storage(Box::new(error)) + } +} +impl From for DataFusionError { + fn from(error: datafusion_storage::path::Error) -> Self { + Self::External(Box::new(error)) + } +} + #[cfg(test)] mod test { use super::*; diff --git a/datafusion/core/Cargo.toml b/datafusion/core/Cargo.toml index 222c0ec688b78..dffaa0a40fa33 100644 --- a/datafusion/core/Cargo.toml +++ b/datafusion/core/Cargo.toml @@ -39,6 +39,10 @@ all-features = true workspace = true [features] +# Default local storage and re-exports for the ObjectStore adapter. +object_store = ["dep:datafusion-storage-object-store"] +# Direct OpenDAL storage adapter; provider services are selected by the embedding application. +opendal = ["dep:datafusion-storage-opendal"] # Enables the `print_config_docs`, `print_runtime_config_docs` and # `print_functions_docs` binaries, which regenerate the docs under # `docs/source/user-guide`. Off by default: they are only run by @@ -63,6 +67,8 @@ compression = [ crypto_expressions = ["datafusion-functions/crypto_expressions"] datetime_expressions = ["datafusion-functions/datetime_expressions"] default = [ + "object_store", + "nested_expressions", "crypto_expressions", "datetime_expressions", @@ -118,6 +124,10 @@ unicode_expressions = [ extended_tests = [] [dependencies] +bytes = { workspace = true } +datafusion-storage-object-store = { workspace = true, optional = true } +datafusion-storage-opendal = { workspace = true, optional = true } +datafusion-storage = { workspace = true } arrow = { workspace = true } arrow-schema = { workspace = true, features = ["canonical_extension_types"] } async-trait = { workspace = true } @@ -125,7 +135,7 @@ bzip2 = { workspace = true, optional = true } chrono = { workspace = true } datafusion-catalog = { workspace = true } datafusion-catalog-listing = { workspace = true } -datafusion-common = { workspace = true, features = ["object_store"] } +datafusion-common = { workspace = true } datafusion-common-runtime = { workspace = true } datafusion-datasource = { workspace = true } datafusion-datasource-arrow = { workspace = true } @@ -155,7 +165,6 @@ indexmap = { workspace = true } itertools = { workspace = true } liblzma = { workspace = true, optional = true } log = { workspace = true } -object_store = { workspace = true } parking_lot = { workspace = true } parquet = { workspace = true, optional = true, default-features = true } serde = { version = "1.0", default-features = false, features = ["derive"], optional = true } @@ -167,6 +176,8 @@ uuid = { workspace = true, features = ["v4", "js"] } zstd = { workspace = true, optional = true } [dev-dependencies] +object_store = { workspace = true } +opendal = { workspace = true, features = ["services-memory"] } async-trait = { workspace = true } criterion = { workspace = true, features = ["async_tokio", "async_futures"] } ctor = { workspace = true } @@ -176,7 +187,6 @@ datafusion-functions-window-common = { workspace = true } datafusion-macros = { workspace = true } datafusion-physical-optimizer = { workspace = true } doc-comment = { workspace = true } -bytes = { workspace = true } env_logger = { workspace = true } glob = { workspace = true } insta = { workspace = true } diff --git a/datafusion/core/benches/push_down_filter.rs b/datafusion/core/benches/push_down_filter.rs index d41085907dbc8..2575f7c1bf74a 100644 --- a/datafusion/core/benches/push_down_filter.rs +++ b/datafusion/core/benches/push_down_filter.rs @@ -21,10 +21,10 @@ use bytes::{BufMut, BytesMut}; use criterion::{Criterion, criterion_group, criterion_main}; use datafusion::config::ConfigOptions; use datafusion::prelude::{ParquetReadOptions, SessionContext}; -use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_physical_optimizer::PhysicalOptimizerRule; use datafusion_physical_optimizer::filter_pushdown::FilterPushdown; use datafusion_physical_plan::ExecutionPlan; +use datafusion_storage::StorageUrl; use object_store::memory::InMemory; use object_store::path::Path; use object_store::{ObjectStore, ObjectStoreExt}; @@ -53,10 +53,13 @@ async fn create_plan() -> Arc { .put(&Path::from("test.parquet"), data.into()) .await .unwrap(); - ctx.register_object_store( - ObjectStoreUrl::parse("memory://").unwrap().as_ref(), - store, - ); + ctx.register_storage( + StorageUrl::parse("memory://").unwrap().as_ref(), + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + store, + )), + ) + .unwrap(); ctx.register_parquet("t", "memory:///", ParquetReadOptions::default()) .await diff --git a/datafusion/core/benches/sql_query_with_io.rs b/datafusion/core/benches/sql_query_with_io.rs index c6600e197374b..94947b0091508 100644 --- a/datafusion/core/benches/sql_query_with_io.rs +++ b/datafusion/core/benches/sql_query_with_io.rs @@ -117,7 +117,13 @@ where async fn setup_context(object_store: Arc) -> SessionContext { let config = SessionConfig::new().with_target_partitions(THREADS); let rt = Arc::new(RuntimeEnv::default()); - rt.register_object_store(&Url::parse("data://my_store").unwrap(), object_store); + rt.register_storage( + &Url::parse("data://my_store").unwrap(), + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + object_store, + )), + ) + .unwrap(); let context = SessionContext::new_with_config_rt(config, rt); for table_id in 0..TABLES { diff --git a/datafusion/core/src/dataframe/mod.rs b/datafusion/core/src/dataframe/mod.rs index ed3dc5ea838b9..e6bda4d7a2a34 100644 --- a/datafusion/core/src/dataframe/mod.rs +++ b/datafusion/core/src/dataframe/mod.rs @@ -42,7 +42,8 @@ use crate::physical_plan::{ }; use crate::prelude::SessionContext; use std::borrow::Cow; -use std::collections::{HashMap, HashSet}; +use std::collections::HashMap; +use std::collections::HashSet; use std::sync::Arc; use arrow::array::{Array, ArrayRef, Int64Array, StringArray}; @@ -51,9 +52,10 @@ use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use arrow::util::display::{ArrayFormatter, FormatOptions}; use arrow_schema::FieldRef; use datafusion_common::config::{CsvOptions, JsonOptions}; +use datafusion_common::not_impl_err; use datafusion_common::{ Column, DFSchema, DataFusionError, ParamValues, ScalarValue, SchemaError, - TableReference, UnnestOptions, exec_err, internal_datafusion_err, not_impl_err, + TableReference, UnnestOptions, exec_err, internal_datafusion_err, plan_datafusion_err, plan_err, unqualified_field_not_found, }; use datafusion_expr::select_expr::SelectExpr; diff --git a/datafusion/core/src/dataframe/parquet.rs b/datafusion/core/src/dataframe/parquet.rs index 1685dff23dff1..a0861cdeca19e 100644 --- a/datafusion/core/src/dataframe/parquet.rs +++ b/datafusion/core/src/dataframe/parquet.rs @@ -184,10 +184,17 @@ mod tests { for compression in test_compressions.into_iter() { let df = test_df.clone(); let tmp_dir = TempDir::new()?; - let local = Arc::new(LocalFileSystem::new_with_prefix(&tmp_dir)?); + let local = Arc::new(LocalFileSystem::new_with_prefix(&tmp_dir).unwrap()); let local_url = Url::parse("file://local").unwrap(); let ctx = &test_df.session_state; - ctx.runtime_env().register_object_store(&local_url, local); + ctx.runtime_env() + .register_storage( + &local_url, + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + local, + )), + ) + .unwrap(); let mut options = TableParquetOptions::default(); options.global.compression = Some(compression.to_string()); df.write_parquet( @@ -233,10 +240,17 @@ mod tests { for rg_size in 1..10 { let df = test_df.clone(); let tmp_dir = TempDir::new()?; - let local = Arc::new(LocalFileSystem::new_with_prefix(&tmp_dir)?); + let local = Arc::new(LocalFileSystem::new_with_prefix(&tmp_dir).unwrap()); let local_url = Url::parse("file://local").unwrap(); let ctx = &test_df.session_state; - ctx.runtime_env().register_object_store(&local_url, local); + ctx.runtime_env() + .register_storage( + &local_url, + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + local, + )), + ) + .unwrap(); let mut options = TableParquetOptions::default(); options.global.max_row_group_size = rg_size; options.global.allow_single_file_parallelism = true; diff --git a/datafusion/core/src/datasource/file_format/csv.rs b/datafusion/core/src/datasource/file_format/csv.rs index 2fb64fd6486e6..8195ed8548790 100644 --- a/datafusion/core/src/datasource/file_format/csv.rs +++ b/datafusion/core/src/datasource/file_format/csv.rs @@ -381,9 +381,10 @@ mod tests { let variable_object_store = Arc::new(VariableStream::new(Bytes::from("1,2,3,4,5\n"), 200)); let object_meta = ObjectMeta { - location: Path::parse("/")?, + location: Path::parse("/").unwrap(), last_modified: DateTime::default(), - size: u64::MAX, + size: variable_object_store.bytes_to_repeat.len() as u64 + * variable_object_store.max_iterations, e_tag: None, version: None, }; @@ -395,8 +396,8 @@ mod tests { let inferred_schema = csv_format .infer_schema( &state, - &(variable_object_store.clone() as Arc), - &[object_meta], + &test_utils::storage::object_store(variable_object_store.clone()), + &[datafusion_storage_object_store::file_info(object_meta)], ) .await?; @@ -439,9 +440,10 @@ mod tests { 1, )); let object_meta = ObjectMeta { - location: Path::parse("/")?, + location: Path::parse("/").unwrap(), last_modified: DateTime::default(), - size: u64::MAX, + size: variable_object_store.bytes_to_repeat.len() as u64 + * variable_object_store.max_iterations, e_tag: None, version: None, }; @@ -456,8 +458,8 @@ mod tests { let inferred_schema = csv_format .infer_schema( &state, - &(variable_object_store.clone() as Arc), - &[object_meta], + &test_utils::storage::object_store(variable_object_store.clone()), + &[datafusion_storage_object_store::file_info(object_meta)], ) .await?; @@ -500,9 +502,9 @@ mod tests { 1, )); let object_meta = ObjectMeta { - location: Path::parse("/")?, + location: Path::parse("/").unwrap(), last_modified: DateTime::default(), - size: u64::MAX, + size: b"c1,c2,c3\n1,1.0,\n,,\n".len() as u64, e_tag: None, version: None, }; @@ -511,8 +513,8 @@ mod tests { let inferred_schema = csv_format .infer_schema( &state, - &(chunked_object_store as Arc), - &[object_meta], + &test_utils::storage::object_store(chunked_object_store), + &[datafusion_storage_object_store::file_info(object_meta)], ) .await?; @@ -559,14 +561,14 @@ mod tests { .schema_infer_max_rec .unwrap_or(DEFAULT_SCHEMA_INFER_MAX_RECORD); let store = Arc::new(integration) as Arc; - let original_stream = store.get(&path).await?; + let original_stream = store.get(&path).await.unwrap(); //convert original_stream to compressed_stream for next step let compressed_stream = file_compression_type.to_owned().convert_to_compress_stream( original_stream .into_stream() - .map_err(DataFusionError::from) + .map_err(|e| DataFusionError::External(Box::new(e))) .boxed(), ); @@ -1375,9 +1377,10 @@ mod tests { let csv_data = Bytes::from("a,b,c\n1,2\n3,4,5\n"); let variable_object_store = Arc::new(VariableStream::new(csv_data, 1)); let object_meta = ObjectMeta { - location: Path::parse("/")?, + location: Path::parse("/").unwrap(), last_modified: DateTime::default(), - size: u64::MAX, + size: variable_object_store.bytes_to_repeat.len() as u64 + * variable_object_store.max_iterations, e_tag: None, version: None, }; @@ -1392,8 +1395,8 @@ mod tests { let inferred_schema = csv_format .infer_schema( &state, - &(variable_object_store.clone() as Arc), - &[object_meta], + &test_utils::storage::object_store(variable_object_store.clone()), + &[datafusion_storage_object_store::file_info(object_meta)], ) .await?; @@ -1452,9 +1455,10 @@ mod tests { let csv_data = Bytes::from("id,a,b,c\n1,foo,bar\n2,foo,bar,baz\n"); let variable_object_store = Arc::new(VariableStream::new(csv_data, 1)); let object_meta = ObjectMeta { - location: Path::parse("/")?, + location: Path::parse("/").unwrap(), last_modified: DateTime::default(), - size: u64::MAX, + size: variable_object_store.bytes_to_repeat.len() as u64 + * variable_object_store.max_iterations, e_tag: None, version: None, }; @@ -1467,8 +1471,8 @@ mod tests { let res = csv_format .infer_schema( &state, - &(variable_object_store.clone() as Arc), - &[object_meta], + &test_utils::storage::object_store(variable_object_store.clone()), + &[datafusion_storage_object_store::file_info(object_meta)], ) .await; @@ -1627,7 +1631,7 @@ mod tests { let path = directory.path().join("duplicate_header.csv"); std::fs::write(&path, "id,value,value\n1,10,100\n")?; - let store = Arc::new(LocalFileSystem::new()) as _; + let store = test_utils::storage::local(); let meta = crate::test::object_store::local_unpartitioned_file(&path); let ctx = SessionContext::new().state(); diff --git a/datafusion/core/src/datasource/file_format/json.rs b/datafusion/core/src/datasource/file_format/json.rs index 1f6f27242e723..261902b8969e7 100644 --- a/datafusion/core/src/datasource/file_format/json.rs +++ b/datafusion/core/src/datasource/file_format/json.rs @@ -52,7 +52,6 @@ mod tests { use datafusion_datasource::file_compression_type::FileCompressionType; use futures::StreamExt; use insta::assert_snapshot; - use object_store::local::LocalFileSystem; use regex::Regex; use rstest::rstest; // ==================== Test Helpers ==================== @@ -72,7 +71,7 @@ mod tests { let (_tmp_dir, path) = create_temp_json(content); let session = SessionContext::new(); let ctx = session.state(); - let store = Arc::new(LocalFileSystem::new()) as _; + let store = test_utils::storage::local(); let format = JsonFormat::default().with_newline_delimited(false); format .infer_schema(&ctx, &store, &[local_unpartitioned_file(&path)]) @@ -209,7 +208,7 @@ mod tests { async fn infer_schema_with_limit() { let session = SessionContext::new(); let ctx = session.state(); - let store = Arc::new(LocalFileSystem::new()) as _; + let store = test_utils::storage::local(); let filename = "tests/data/schema_infer_limit.json"; let format = JsonFormat::default().with_schema_infer_max_rec(3); diff --git a/datafusion/core/src/datasource/file_format/mod.rs b/datafusion/core/src/datasource/file_format/mod.rs index c46b472bd6404..72bf8b1b45917 100644 --- a/datafusion/core/src/datasource/file_format/mod.rs +++ b/datafusion/core/src/datasource/file_format/mod.rs @@ -42,7 +42,7 @@ pub(crate) mod test_util { use datafusion_datasource::TableSchema; use datafusion_datasource::file_scan_config::FileScanConfigBuilder; use datafusion_datasource::{PartitionedFile, file_format::FileFormat}; - use datafusion_execution::object_store::ObjectStoreUrl; + use datafusion_storage::StorageUrl; use std::sync::Arc; use crate::test::object_store::local_unpartitioned_file; @@ -56,7 +56,7 @@ pub(crate) mod test_util { projection: Option>, limit: Option, ) -> Result> { - let store = Arc::new(object_store::local::LocalFileSystem::new()) as _; + let store = test_utils::storage::local(); let meta = local_unpartitioned_file(format!("{store_root}/{file_name}")); let file_schema = if let Some(file_schema) = schema { @@ -79,7 +79,7 @@ pub(crate) mod test_util { .create_physical_plan( state, FileScanConfigBuilder::new( - ObjectStoreUrl::local_filesystem(), + StorageUrl::local_filesystem(), format.file_source(table_schema), ) .with_file_groups(file_groups) @@ -112,9 +112,15 @@ mod tests { let ctx = SessionContext::new(); // register a local file system object store for /tmp directory let tmp_dir = TempDir::new()?; - let local = Arc::new(LocalFileSystem::new_with_prefix(&tmp_dir)?); + let local = Arc::new(LocalFileSystem::new_with_prefix(&tmp_dir).unwrap()); let local_url = Url::parse("file://local").unwrap(); - ctx.register_object_store(&local_url, local); + ctx.register_storage( + &local_url, + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + local, + )), + ) + .unwrap(); let options = CsvReadOptions::default() .schema_infer_max_records(2) diff --git a/datafusion/core/src/datasource/file_format/parquet.rs b/datafusion/core/src/datasource/file_format/parquet.rs index bfcfb74848861..9ca568a1c5624 100644 --- a/datafusion/core/src/datasource/file_format/parquet.rs +++ b/datafusion/core/src/datasource/file_format/parquet.rs @@ -23,7 +23,7 @@ pub use datafusion_datasource_parquet::file_format::*; pub(crate) mod test_util { use arrow::array::RecordBatch; use datafusion_common::Result; - use object_store::ObjectMeta; + use datafusion_storage::FileInfo; use crate::test::object_store::local_unpartitioned_file; @@ -38,7 +38,7 @@ pub(crate) mod test_util { pub async fn store_parquet( batches: Vec, multi_page: bool, - ) -> Result<(Vec, Vec)> { + ) -> Result<(Vec, Vec)> { /// How many rows per page should be written const ROWS_PER_PAGE: usize = 2; /// write batches chunk_size rows at a time @@ -138,12 +138,12 @@ mod tests { ParquetFormat, ParquetFormatFactory, ParquetSink, }; use datafusion_execution::TaskContext; - use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_execution::runtime_env::RuntimeEnv; use datafusion_expr::dml::InsertOp; use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext}; use datafusion_physical_plan::stream::RecordBatchStreamAdapter; use datafusion_physical_plan::{ExecutionPlan, collect}; + use datafusion_storage::StorageUrl; use crate::test_util::bounded_stream; use arrow::array::{ @@ -186,7 +186,7 @@ mod tests { let batch1 = RecordBatch::try_from_iter(vec![("c1", c1.clone())])?; let batch2 = RecordBatch::try_from_iter(vec![("c2", c2)])?; - let store = Arc::new(LocalFileSystem::new()) as _; + let store = test_utils::storage::local(); let (meta, _files) = store_parquet(vec![batch1, batch2], false).await?; let session = SessionContext::new(); @@ -200,8 +200,12 @@ mod tests { let file_metadata_cache = ctx.runtime_env().cache_manager.get_file_metadata_cache(); - let stats = DFParquetMetadata::new(&store, &meta[0]) - .with_file_metadata_cache(Some(Arc::clone(&file_metadata_cache))) + let file_info = meta[0].clone(); + let reader = store + .open(&file_info, datafusion_storage::FileAccessContext::default()) + .await?; + let stats = DFParquetMetadata::new(reader.as_ref(), &file_info) + .with_file_metadata_cache(Some(Arc::clone(&file_metadata_cache)), store.id()) .fetch_statistics(&schema) .await?; @@ -211,8 +215,12 @@ mod tests { assert_eq!(c1_stats.null_count, Precision::Exact(1)); assert_eq!(c2_stats.null_count, Precision::Exact(3)); - let stats = DFParquetMetadata::new(&store, &meta[1]) - .with_file_metadata_cache(Some(Arc::clone(&file_metadata_cache))) + let file_info = meta[1].clone(); + let reader = store + .open(&file_info, datafusion_storage::FileAccessContext::default()) + .await?; + let stats = DFParquetMetadata::new(reader.as_ref(), &file_info) + .with_file_metadata_cache(Some(Arc::clone(&file_metadata_cache)), store.id()) .fetch_statistics(&schema) .await?; @@ -253,7 +261,7 @@ mod tests { let batch2 = RecordBatch::try_from_iter(vec![("c", c2.clone()), ("d", c2.clone())])?; - let store = Arc::new(LocalFileSystem::new()) as _; + let store = test_utils::storage::local(); let (meta, _files) = store_parquet(vec![batch1, batch2], false).await?; let session = SessionContext::new(); @@ -375,6 +383,7 @@ mod tests { let store = Arc::new(RequestCountingObjectStore::new(Arc::new( LocalFileSystem::new(), ))); + let binding = test_utils::storage::object_store(store.upcast()); let (meta, _files) = store_parquet(vec![batch1, batch2], false).await?; let session = SessionContext::new(); @@ -384,13 +393,19 @@ mod tests { // for the remaining metadata let file_metadata_cache = ctx.runtime_env().cache_manager.get_file_metadata_cache(); - let df_meta = DFParquetMetadata::new(store.as_ref(), &meta[0]) + let file_info = meta[0].clone(); + let reader = binding + .open(&file_info, datafusion_storage::FileAccessContext::default()) + .await?; + let df_meta = DFParquetMetadata::new(reader.as_ref(), &file_info) .with_metadata_size_hint(Some(9)); df_meta.fetch_metadata().await?; assert_eq!(store.request_count(), 2); - let df_meta = - df_meta.with_file_metadata_cache(Some(Arc::clone(&file_metadata_cache))); + let df_meta = df_meta.with_file_metadata_cache( + Some(Arc::clone(&file_metadata_cache)), + binding.id(), + ); // Increases by 3 because cache has no entries yet df_meta.fetch_metadata().await?; @@ -401,7 +416,7 @@ mod tests { assert_eq!(store.request_count(), 5); // Increase by 2 because `get_file_metadata_cache()` is None - let df_meta = df_meta.with_file_metadata_cache(None); + let df_meta = df_meta.with_file_metadata_cache(None, binding.id()); df_meta.fetch_metadata().await?; assert_eq!(store.request_count(), 7); @@ -413,15 +428,17 @@ mod tests { .with_metadata_size_hint(Some(9)) .with_force_view_types(force_views); // Increase by 3, partial cache being used. - let _schema = format.infer_schema(&ctx, &store.upcast(), &meta).await?; + let _schema = format.infer_schema(&ctx, &binding, &meta).await?; assert_eq!(store.request_count(), 10); // No increase, full cache being used. - let schema = format.infer_schema(&ctx, &store.upcast(), &meta).await?; + let schema = format.infer_schema(&ctx, &binding, &meta).await?; assert_eq!(store.request_count(), 10); // No increase, cache being used - let df_meta = - df_meta.with_file_metadata_cache(Some(Arc::clone(&file_metadata_cache))); + let df_meta = df_meta.with_file_metadata_cache( + Some(Arc::clone(&file_metadata_cache)), + binding.id(), + ); let stats = df_meta.fetch_statistics(&schema).await?; assert_eq!(store.request_count(), 10); @@ -434,10 +451,15 @@ mod tests { let store = Arc::new(RequestCountingObjectStore::new(Arc::new( LocalFileSystem::new(), ))); + let binding = test_utils::storage::object_store(store.upcast()); // Use the file size as the hint so we can get the full metadata from the first fetch let size_hint = meta[0].size as usize; - let df_meta = DFParquetMetadata::new(store.as_ref(), &meta[0]) + let file_info = meta[0].clone(); + let reader = binding + .open(&file_info, datafusion_storage::FileAccessContext::default()) + .await?; + let df_meta = DFParquetMetadata::new(reader.as_ref(), &file_info) .with_metadata_size_hint(Some(size_hint)); df_meta.fetch_metadata().await?; @@ -448,8 +470,10 @@ mod tests { let ctx = session.state(); let file_metadata_cache = ctx.runtime_env().cache_manager.get_file_metadata_cache(); - let df_meta = - df_meta.with_file_metadata_cache(Some(Arc::clone(&file_metadata_cache))); + let df_meta = df_meta.with_file_metadata_cache( + Some(Arc::clone(&file_metadata_cache)), + binding.id(), + ); // Increases by 1 because cache has no entries yet and new session context df_meta.fetch_metadata().await?; assert_eq!(store.request_count(), 2); @@ -459,7 +483,7 @@ mod tests { assert_eq!(store.request_count(), 2); // Increase by 1 because `get_file_metadata_cache` is None - let df_meta = df_meta.with_file_metadata_cache(None); + let df_meta = df_meta.with_file_metadata_cache(None, binding.id()); df_meta.fetch_metadata().await?; assert_eq!(store.request_count(), 3); @@ -467,14 +491,16 @@ mod tests { .with_metadata_size_hint(Some(size_hint)) .with_force_view_types(force_views); // Increase by 1, partial cache being used. - let _schema = format.infer_schema(&ctx, &store.upcast(), &meta).await?; + let _schema = format.infer_schema(&ctx, &binding, &meta).await?; assert_eq!(store.request_count(), 4); // No increase, full cache being used. - let schema = format.infer_schema(&ctx, &store.upcast(), &meta).await?; + let schema = format.infer_schema(&ctx, &binding, &meta).await?; assert_eq!(store.request_count(), 4); // No increase, cache being used - let df_meta = - df_meta.with_file_metadata_cache(Some(Arc::clone(&file_metadata_cache))); + let df_meta = df_meta.with_file_metadata_cache( + Some(Arc::clone(&file_metadata_cache)), + binding.id(), + ); let stats = df_meta.fetch_statistics(&schema).await?; assert_eq!(store.request_count(), 4); @@ -487,20 +513,27 @@ mod tests { let store = Arc::new(RequestCountingObjectStore::new(Arc::new( LocalFileSystem::new(), ))); + let binding = test_utils::storage::object_store(store.upcast()); // Use a size hint larger than the file size to make sure we don't panic let size_hint = (meta[0].size + 100) as usize; - let df_meta = DFParquetMetadata::new(store.as_ref(), &meta[0]) + let file_info = meta[0].clone(); + let reader = binding + .open(&file_info, datafusion_storage::FileAccessContext::default()) + .await?; + let df_meta = DFParquetMetadata::new(reader.as_ref(), &file_info) .with_metadata_size_hint(Some(size_hint)); df_meta.fetch_metadata().await?; assert_eq!(store.request_count(), 1); - // No increase because cache has an entry - let df_meta = - df_meta.with_file_metadata_cache(Some(Arc::clone(&file_metadata_cache))); + // A new binding must not reuse entries from the previous registration. + let df_meta = df_meta.with_file_metadata_cache( + Some(Arc::clone(&file_metadata_cache)), + binding.id(), + ); df_meta.fetch_metadata().await?; - assert_eq!(store.request_count(), 1); + assert_eq!(store.request_count(), 2); Ok(()) } @@ -541,22 +574,30 @@ mod tests { let store = Arc::new(RequestCountingObjectStore::new(Arc::new( LocalFileSystem::new(), ))); + let binding = test_utils::storage::object_store(store.upcast()); let (files, _file_names) = store_parquet(vec![batch1], false).await?; let state = SessionContext::new().state(); // Make metadata size hint None to keep original behavior let format = ParquetFormat::default().with_metadata_size_hint(None); - let _schema = format.infer_schema(&state, &store.upcast(), &files).await?; + let _schema = format.infer_schema(&state, &binding, &files).await?; assert_eq!(store.request_count(), 3); // No increase, cache being used. - let schema = format.infer_schema(&state, &store.upcast(), &files).await?; + let schema = format.infer_schema(&state, &binding, &files).await?; assert_eq!(store.request_count(), 3); // No increase in request count because cache is not empty let file_metadata_cache = state.runtime_env().cache_manager.get_file_metadata_cache(); - let stats = DFParquetMetadata::new(store.as_ref(), &files[0]) - .with_file_metadata_cache(Some(Arc::clone(&file_metadata_cache))) + let file_info = files[0].clone(); + let reader = binding + .open(&file_info, datafusion_storage::FileAccessContext::default()) + .await?; + let stats = DFParquetMetadata::new(reader.as_ref(), &file_info) + .with_file_metadata_cache( + Some(Arc::clone(&file_metadata_cache)), + binding.id(), + ) .fetch_statistics(&schema) .await?; assert_eq!(stats.num_rows, Precision::Exact(4)); @@ -610,6 +651,7 @@ mod tests { let store = Arc::new(RequestCountingObjectStore::new(Arc::new( LocalFileSystem::new(), ))); + let binding = test_utils::storage::object_store(store.upcast()); let (files, _file_names) = store_parquet(vec![batch1, batch2], false).await?; let force_views = match force_views { @@ -622,7 +664,7 @@ mod tests { let format = ParquetFormat::default() .with_force_view_types(force_views) .with_metadata_size_hint(None); - let schema = format.infer_schema(&state, &store.upcast(), &files).await?; + let schema = format.infer_schema(&state, &binding, &files).await?; assert_eq!(store.request_count(), 6); let null_i64 = ScalarValue::Int64(None); @@ -635,8 +677,15 @@ mod tests { // No increase in request count because cache is not empty let file_metadata_cache = state.runtime_env().cache_manager.get_file_metadata_cache(); - let stats = DFParquetMetadata::new(store.as_ref(), &files[0]) - .with_file_metadata_cache(Some(Arc::clone(&file_metadata_cache))) + let file_info = files[0].clone(); + let reader = binding + .open(&file_info, datafusion_storage::FileAccessContext::default()) + .await?; + let stats = DFParquetMetadata::new(reader.as_ref(), &file_info) + .with_file_metadata_cache( + Some(Arc::clone(&file_metadata_cache)), + binding.id(), + ) .fetch_statistics(&schema) .await?; assert_eq!(store.request_count(), 6); @@ -664,8 +713,12 @@ mod tests { assert_eq!(c2_stats.min_value, Precision::Exact(null_i64.clone())); // No increase in request count because cache is not empty - let stats = DFParquetMetadata::new(store.as_ref(), &files[1]) - .with_file_metadata_cache(Some(Arc::clone(&file_metadata_cache))) + let file_info = files[1].clone(); + let stats = DFParquetMetadata::new(reader.as_ref(), &file_info) + .with_file_metadata_cache( + Some(Arc::clone(&file_metadata_cache)), + binding.id(), + ) .fetch_statistics(&schema) .await?; assert_eq!(store.request_count(), 6); @@ -1267,8 +1320,13 @@ mod tests { let runtime = RuntimeEnv::default(); runtime - .object_store_registry - .register_store(store_url, local); + .register_storage( + store_url, + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + local, + )), + ) + .unwrap(); Arc::new( TaskContext::default() @@ -1542,9 +1600,10 @@ mod tests { let field_a = Field::new("a", DataType::Utf8, false); let field_b = Field::new("b", DataType::Utf8, false); let schema = Arc::new(Schema::new(vec![field_a, field_b])); - let object_store_url = ObjectStoreUrl::local_filesystem(); + let object_store_url = StorageUrl::local_filesystem(); let file_sink_config = FileSinkConfig { + storage: None, original_url: String::default(), object_store_url: object_store_url.clone(), file_group: FileGroup::new(vec![PartitionedFile::new("/tmp".to_string(), 1)]), @@ -1598,7 +1657,7 @@ mod tests { ); let (path, parquet_meta_data) = written.take(1).next().unwrap(); - Ok((path, parquet_meta_data)) + Ok((Path::from(path.as_ref()), parquet_meta_data)) } fn assert_file_metadata( @@ -1633,10 +1692,11 @@ mod tests { let field_a = Field::new("a", DataType::Utf8, false); let field_b = Field::new("b", DataType::Utf8, false); let schema = Arc::new(Schema::new(vec![field_a, field_b])); - let object_store_url = ObjectStoreUrl::local_filesystem(); + let object_store_url = StorageUrl::local_filesystem(); // set file config to include partitioning on field_a let file_sink_config = FileSinkConfig { + storage: None, original_url: String::default(), object_store_url: object_store_url.clone(), file_group: FileGroup::new(vec![PartitionedFile::new("/tmp".to_string(), 1)]), @@ -1722,9 +1782,10 @@ mod tests { let field_a = Field::new("a", DataType::Utf8, false); let field_b = Field::new("b", DataType::Utf8, false); let schema = Arc::new(Schema::new(vec![field_a, field_b])); - let object_store_url = ObjectStoreUrl::local_filesystem(); + let object_store_url = StorageUrl::local_filesystem(); let file_sink_config = FileSinkConfig { + storage: None, original_url: String::default(), object_store_url: object_store_url.clone(), file_group: FileGroup::new(vec![PartitionedFile::new( @@ -1839,7 +1900,7 @@ mod tests { ], )?; - let store = Arc::new(LocalFileSystem::new()) as _; + let store = test_utils::storage::local(); let (meta, _files) = store_parquet(vec![batch], false).await?; let ctx = SessionContext::new().state(); diff --git a/datafusion/core/src/datasource/listing/table.rs b/datafusion/core/src/datasource/listing/table.rs index 982766dc88519..2d4f19456ac11 100644 --- a/datafusion/core/src/datasource/listing/table.rs +++ b/datafusion/core/src/datasource/listing/table.rs @@ -85,16 +85,20 @@ fn infer_boxed( /// Body of [`ListingTableConfigExt::infer_options`]. fn infer_options_boxed( - config: ListingTableConfig, + mut config: ListingTableConfig, state: &dyn Session, ) -> BoxFuture<'_, datafusion_common::Result> { Box::pin(async move { let store = if let Some(url) = config.table_paths.first() { - state.runtime_env().object_store(url)? + match &config.storage { + Some(storage) => std::sync::Arc::clone(storage), + None => state.runtime_env().storage(url)?, + } } else { return Ok(config); }; + config.storage = Some(std::sync::Arc::clone(&store)); let file = config .table_paths .first() @@ -582,7 +586,8 @@ mod tests { let head_concurrency_store = ensure_head_concurrency(store, expected_concurrency); let url = Url::parse("test://").unwrap(); - ctx.register_object_store(&url, head_concurrency_store.clone()); + ctx.register_storage(&url, Arc::clone(head_concurrency_store.storage())) + .unwrap(); let format = JsonFormat::default(); diff --git a/datafusion/core/src/datasource/listing_table_factory.rs b/datafusion/core/src/datasource/listing_table_factory.rs index 1e597e38fb5b1..4d223a11415cd 100644 --- a/datafusion/core/src/datasource/listing_table_factory.rs +++ b/datafusion/core/src/datasource/listing_table_factory.rs @@ -126,16 +126,16 @@ impl ListingTableFactory { // object stores is intentionally not supported (see // https://github.com/apache/datafusion/issues/16303); reject it here with // a clear error rather than producing incorrect results at scan time. - let object_store_url = first_path.object_store(); + let object_store_url = first_path.storage_url(); if let Some(other) = table_paths .iter() - .find(|path| path.object_store() != object_store_url) + .find(|path| path.storage_url() != object_store_url) { return plan_err!( "All locations of a CREATE EXTERNAL TABLE must be on the same \ object store, but found '{}' and '{}'", object_store_url.as_str(), - other.object_store().as_str() + other.storage_url().as_str() ); } diff --git a/datafusion/core/src/datasource/mod.rs b/datafusion/core/src/datasource/mod.rs index de54078aafef4..d687fb95008e8 100644 --- a/datafusion/core/src/datasource/mod.rs +++ b/datafusion/core/src/datasource/mod.rs @@ -24,6 +24,7 @@ pub mod file_format; pub mod listing; pub mod listing_table_factory; mod memory_test; + pub mod physical_plan; pub mod provider; mod view_test; @@ -47,14 +48,13 @@ pub use datafusion_datasource::schema_adapter; pub use datafusion_datasource::sink; pub use datafusion_datasource::source; pub use datafusion_datasource::table_schema; -pub use datafusion_execution::object_store; + pub use datafusion_physical_expr::create_ordering; #[cfg(all(test, feature = "parquet"))] mod tests { use crate::prelude::SessionContext; - use ::object_store::{ObjectMeta, path::Path}; use arrow::{ array::Int32Array, datatypes::{DataType, Field, Schema, SchemaRef}, @@ -75,6 +75,7 @@ mod tests { }; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use datafusion_physical_plan::collect; + use datafusion_storage::{FileInfo, path::Path}; use std::{fs, sync::Arc}; use tempfile::TempDir; use url::Url; @@ -85,7 +86,7 @@ mod tests { // record batches returned from parquet. This can be useful for schema evolution // where older files may not have all columns. - use datafusion_execution::object_store::ObjectStoreUrl; + use datafusion_storage::StorageUrl; let tmp_dir = TempDir::new().unwrap(); let table_dir = tmp_dir.path().join("parquet_test"); fs::DirBuilder::new().create(table_dir.as_path()).unwrap(); @@ -108,7 +109,7 @@ mod tests { let url = Url::from_file_path(path.canonicalize().unwrap()).unwrap(); let location = Path::from_url_path(url.path()).unwrap(); let metadata = fs::metadata(path.as_path()).expect("Local file metadata"); - let meta = ObjectMeta { + let meta = FileInfo { location, last_modified: metadata.modified().map(chrono::DateTime::from).unwrap(), size: metadata.len(), @@ -124,7 +125,7 @@ mod tests { let schema = Arc::new(Schema::new(vec![f1.clone(), f2.clone()])); let source = Arc::new(ParquetSource::new(Arc::clone(&schema))); let base_conf = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source) + FileScanConfigBuilder::new(StorageUrl::local_filesystem(), source) .with_file(partitioned_file) .with_expr_adapter(Some(Arc::new(TestPhysicalExprAdapterFactory))) .build(); diff --git a/datafusion/core/src/datasource/physical_plan/avro.rs b/datafusion/core/src/datasource/physical_plan/avro.rs index c9ee2cc407783..a48b26bcccb16 100644 --- a/datafusion/core/src/datasource/physical_plan/avro.rs +++ b/datafusion/core/src/datasource/physical_plan/avro.rs @@ -37,8 +37,8 @@ mod tests { use datafusion_datasource::{PartitionedFile, TableSchemaBuilder}; use datafusion_datasource_avro::AvroFormat; use datafusion_datasource_avro::source::AvroSource; - use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_physical_plan::ExecutionPlan; + use datafusion_storage::StorageUrl; use datafusion_datasource::source::DataSourceExec; use futures::StreamExt; @@ -71,18 +71,29 @@ mod tests { let state = session_ctx.state(); let url = Url::parse("file://").unwrap(); - session_ctx.register_object_store(&url, store.clone()); + session_ctx + .register_storage( + &url, + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + store.clone(), + )), + ) + .unwrap(); let testdata = test_util::arrow_test_data(); let filename = format!("{testdata}/avro/alltypes_plain.avro"); let meta = local_unpartitioned_file(filename); let file_schema = AvroFormat {} - .infer_schema(&state, &store, std::slice::from_ref(&meta)) + .infer_schema( + &state, + &test_utils::storage::object_store(store), + std::slice::from_ref(&meta), + ) .await?; let source = Arc::new(AvroSource::new(Arc::clone(&file_schema))); - let conf = FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source) + let conf = FileScanConfigBuilder::new(StorageUrl::local_filesystem(), source) .with_file(meta.into()) .with_projection_indices(Some(vec![0, 1, 2]))? .build(); @@ -139,8 +150,8 @@ mod tests { let testdata = test_util::arrow_test_data(); let filename = format!("{testdata}/avro/alltypes_plain.avro"); - let object_store = Arc::new(LocalFileSystem::new()) as _; - let object_store_url = ObjectStoreUrl::local_filesystem(); + let object_store = test_utils::storage::local(); + let object_store_url = StorageUrl::local_filesystem(); let meta = local_unpartitioned_file(filename); let actual_schema = AvroFormat {} .infer_schema(&state, &object_store, std::slice::from_ref(&meta)) @@ -212,8 +223,8 @@ mod tests { let testdata = test_util::arrow_test_data(); let filename = format!("{testdata}/avro/alltypes_plain.avro"); - let object_store = Arc::new(LocalFileSystem::new()) as _; - let object_store_url = ObjectStoreUrl::local_filesystem(); + let object_store = test_utils::storage::local(); + let object_store_url = StorageUrl::local_filesystem(); let meta = local_unpartitioned_file(filename); let file_schema = AvroFormat {} .infer_schema(&state, &object_store, std::slice::from_ref(&meta)) diff --git a/datafusion/core/src/datasource/physical_plan/csv.rs b/datafusion/core/src/datasource/physical_plan/csv.rs index 7980df87fa576..1801c3c972079 100644 --- a/datafusion/core/src/datasource/physical_plan/csv.rs +++ b/datafusion/core/src/datasource/physical_plan/csv.rs @@ -485,7 +485,13 @@ mod tests { ) -> Result<()> { let ctx = SessionContext::new(); let url = Url::parse("file://").unwrap(); - ctx.register_object_store(&url, store.clone()); + ctx.register_storage( + &url, + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + store.clone(), + )), + ) + .unwrap(); let task_ctx = ctx.task_ctx(); @@ -564,7 +570,14 @@ mod tests { store.put(&path, data.into()).await.unwrap(); let url = Url::parse("memory://").unwrap(); - session_ctx.register_object_store(&url, Arc::new(store)); + session_ctx + .register_storage( + &url, + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + Arc::new(store), + )), + ) + .unwrap(); let df = session_ctx .read_csv("memory:///", CsvReadOptions::new()) @@ -593,7 +606,14 @@ mod tests { store.put(&path, data.into()).await.unwrap(); let url = Url::parse("memory://").unwrap(); - session_ctx.register_object_store(&url, Arc::new(store)); + session_ctx + .register_storage( + &url, + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + Arc::new(store), + )), + ) + .unwrap(); let df = session_ctx .read_csv("memory:///", CsvReadOptions::new().terminator(Some(b'\r'))) @@ -690,9 +710,15 @@ mod tests { // register a local file system object store let tmp_dir = TempDir::new()?; - let local = Arc::new(LocalFileSystem::new_with_prefix(&tmp_dir)?); + let local = Arc::new(LocalFileSystem::new_with_prefix(&tmp_dir).unwrap()); let local_url = Url::parse("file://local").unwrap(); - ctx.register_object_store(&local_url, local); + ctx.register_storage( + &local_url, + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + local, + )), + ) + .unwrap(); let options = CsvReadOptions::default() .schema_infer_max_records(2) .has_header(true); @@ -736,10 +762,16 @@ mod tests { // register a local file system object store let tmp_dir = TempDir::new()?; - let local = Arc::new(LocalFileSystem::new_with_prefix(&tmp_dir)?); + let local = Arc::new(LocalFileSystem::new_with_prefix(&tmp_dir).unwrap()); let local_url = Url::parse("file://local").unwrap(); - ctx.register_object_store(&local_url, local); + ctx.register_storage( + &local_url, + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + local, + )), + ) + .unwrap(); // execute a simple query and write the results to CSV let out_dir = tmp_dir.as_ref().to_str().unwrap().to_string() + "/out/"; diff --git a/datafusion/core/src/datasource/physical_plan/json.rs b/datafusion/core/src/datasource/physical_plan/json.rs index 6b4361e0c4d07..ea170c0df99e4 100644 --- a/datafusion/core/src/datasource/physical_plan/json.rs +++ b/datafusion/core/src/datasource/physical_plan/json.rs @@ -41,8 +41,8 @@ mod tests { use datafusion_datasource::file_format::FileFormat; use datafusion_datasource_json::JsonFormat; use datafusion_execution::config::SessionConfig; - use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_physical_plan::ExecutionPlan; + use datafusion_storage::StorageUrl; use arrow::array::Array; use arrow::datatypes::SchemaRef; @@ -64,9 +64,9 @@ mod tests { state: &SessionState, file_compression_type: FileCompressionType, work_dir: &Path, - ) -> (ObjectStoreUrl, Vec, SchemaRef) { - let store_url = ObjectStoreUrl::local_filesystem(); - let store = state.runtime_env().object_store(&store_url).unwrap(); + ) -> (StorageUrl, Vec, SchemaRef) { + let store_url = StorageUrl::local_filesystem(); + let store = state.runtime_env().storage(&store_url).unwrap(); let filename = "1.json"; let json_format: Arc = Arc::new(JsonFormat::default()); @@ -103,7 +103,13 @@ mod tests { ) -> Result<()> { let ctx = SessionContext::new(); let url = Url::parse("file://").unwrap(); - ctx.register_object_store(&url, store.clone()); + ctx.register_storage( + &url, + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + store.clone(), + )), + ) + .unwrap(); let filename = "1.json"; let tmp_dir = TempDir::new()?; let json_format: Arc = Arc::new(JsonFormat::default()); @@ -127,7 +133,7 @@ mod tests { .location .as_ref(); - let store_url = ObjectStoreUrl::local_filesystem(); + let store_url = StorageUrl::local_filesystem(); let url: &Url = store_url.as_ref(); let path_buf = Path::new(url.path()).join(path); let path = path_buf.to_str().unwrap(); @@ -394,9 +400,15 @@ mod tests { // register a local file system object store for /tmp directory let tmp_dir = TempDir::new()?; - let local = Arc::new(LocalFileSystem::new_with_prefix(&tmp_dir)?); + let local = Arc::new(LocalFileSystem::new_with_prefix(&tmp_dir).unwrap()); let local_url = Url::parse("file://local").unwrap(); - ctx.register_object_store(&local_url, local); + ctx.register_storage( + &local_url, + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + local, + )), + ) + .unwrap(); // execute a simple query and write the results to CSV let out_dir = tmp_dir.as_ref().to_str().unwrap().to_string() + "/out/"; @@ -488,9 +500,15 @@ mod tests { let ctx = SessionContext::new(); // register a local file system object store for /tmp directory let tmp_dir = TempDir::new()?; - let local = Arc::new(LocalFileSystem::new_with_prefix(&tmp_dir)?); + let local = Arc::new(LocalFileSystem::new_with_prefix(&tmp_dir).unwrap()); let local_url = Url::parse("file://local").unwrap(); - ctx.register_object_store(&local_url, local); + ctx.register_storage( + &local_url, + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + local, + )), + ) + .unwrap(); let options = CsvReadOptions::default() .schema_infer_max_records(2) .has_header(true); diff --git a/datafusion/core/src/datasource/physical_plan/mod.rs b/datafusion/core/src/datasource/physical_plan/mod.rs index 8e4855afa66bb..34c19b7048256 100644 --- a/datafusion/core/src/datasource/physical_plan/mod.rs +++ b/datafusion/core/src/datasource/physical_plan/mod.rs @@ -51,3 +51,6 @@ pub use datafusion_datasource::file_sink_config::*; pub use datafusion_datasource::file_stream::{ FileOpenFuture, FileOpener, FileStream, FileStreamBuilder, OnError, }; + +#[cfg(feature = "parquet")] +pub use datafusion_datasource_parquet::storage::StorageParquetTable; diff --git a/datafusion/core/src/datasource/physical_plan/parquet.rs b/datafusion/core/src/datasource/physical_plan/parquet.rs index d562bbe8490f4..059db68529256 100644 --- a/datafusion/core/src/datasource/physical_plan/parquet.rs +++ b/datafusion/core/src/datasource/physical_plan/parquet.rs @@ -59,7 +59,6 @@ mod tests { use datafusion_datasource_parquet::{ DefaultParquetFileReaderFactory, ParquetFileReaderFactory, ParquetFormat, }; - use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_expr::{Expr, col, lit, when}; use datafusion_physical_expr::planner::logical2physical; use datafusion_physical_plan::analyze::AnalyzeExecBuilder; @@ -68,15 +67,16 @@ mod tests { ExecutionPlanMetricsSet, MetricValue, MetricsSet, }; use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties}; + use datafusion_storage::StorageUrl; use chrono::{TimeZone, Utc}; use datafusion_datasource::file_groups::FileGroup; use futures::StreamExt; use insta; use insta::assert_snapshot; + use object_store::ObjectStore; use object_store::local::LocalFileSystem; use object_store::path::Path; - use object_store::{ObjectMeta, ObjectStore}; use parquet::arrow::ArrowWriter; use parquet::file::properties::WriterProperties; use tempfile::TempDir; @@ -195,7 +195,7 @@ mod tests { source: Arc, ) -> Arc { let base_config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source) + FileScanConfigBuilder::new(StorageUrl::local_filesystem(), source) .with_file_group(file_group) .with_projection_indices(self.projection.clone()) .unwrap() @@ -1535,7 +1535,11 @@ mod tests { #[tokio::test] async fn parquet_exec_with_range() -> Result<()> { - fn file_range(meta: &ObjectMeta, start: i64, end: i64) -> PartitionedFile { + fn file_range( + meta: &datafusion_storage::FileInfo, + start: i64, + end: i64, + ) -> PartitionedFile { PartitionedFile::new_from_meta(meta.clone()).with_range(start, end) } @@ -1546,7 +1550,7 @@ mod tests { file_schema: SchemaRef, ) -> Result<()> { let config = FileScanConfigBuilder::new( - ObjectStoreUrl::local_filesystem(), + StorageUrl::local_filesystem(), Arc::new(ParquetSource::new(file_schema)), ) .with_file_groups(file_groups) @@ -1580,7 +1584,7 @@ mod tests { let meta = local_unpartitioned_file(filename); - let store = Arc::new(LocalFileSystem::new()) as _; + let store = test_utils::storage::local(); let file_schema = ParquetFormat::default() .infer_schema(&state, &store, std::slice::from_ref(&meta)) .await?; @@ -1605,8 +1609,8 @@ mod tests { let state = session_ctx.state(); let task_ctx = session_ctx.task_ctx(); - let object_store_url = ObjectStoreUrl::local_filesystem(); - let store = state.runtime_env().object_store(&object_store_url).unwrap(); + let object_store_url = StorageUrl::local_filesystem(); + let store = state.runtime_env().storage(&object_store_url).unwrap(); let testdata = datafusion_common::test_util::parquet_test_data(); let filename = format!("{testdata}/alltypes_plain.parquet"); @@ -1706,17 +1710,18 @@ mod tests { .unwrap() .join("invalid.parquet"); - let partitioned_file = PartitionedFile::new_from_meta(ObjectMeta { - location, - last_modified: Utc.timestamp_nanos(0), - size: 1337, - e_tag: None, - version: None, - }); + let partitioned_file = + PartitionedFile::new_from_meta(datafusion_storage::FileInfo { + location: datafusion_storage::path::Path::from(location.as_ref()), + last_modified: Utc.timestamp_nanos(0), + size: 1337, + e_tag: None, + version: None, + }); let file_schema = Arc::new(Schema::empty()); let config = FileScanConfigBuilder::new( - ObjectStoreUrl::local_filesystem(), + StorageUrl::local_filesystem(), Arc::new(ParquetSource::new(file_schema)), ) .with_file(partitioned_file) @@ -1727,7 +1732,9 @@ mod tests { let mut results = parquet_exec.execute(0, state.task_ctx())?; let batch = results.next().await.unwrap(); // invalid file should produce an error to that effect - assert_contains!(batch.unwrap_err().to_string(), "invalid.parquet not found"); + let error = batch.unwrap_err().to_string(); + assert_contains!(&error, "File not found:"); + assert_contains!(error, "invalid.parquet"); assert!(results.next().await.is_none()); Ok(()) @@ -2207,9 +2214,15 @@ mod tests { .await?; // register a local file system object store for /tmp directory - let local = Arc::new(LocalFileSystem::new_with_prefix(&tmp_dir)?); + let local = Arc::new(LocalFileSystem::new_with_prefix(&tmp_dir).unwrap()); let local_url = Url::parse("file://local").unwrap(); - ctx.register_object_store(&local_url, local); + ctx.register_storage( + &local_url, + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + local, + )), + ) + .unwrap(); // Configure listing options let file_format = ParquetFormat::default().with_enable_pruning(true); @@ -2441,7 +2454,9 @@ mod tests { impl TrackingParquetFileReaderFactory { fn new(store: Arc) -> Self { Self { - inner: Arc::new(DefaultParquetFileReaderFactory::new(store)) as _, + inner: Arc::new(DefaultParquetFileReaderFactory::new( + test_utils::storage::object_store(store), + )) as _, metadata_size_hint_calls: Arc::new(Mutex::new(vec![])), } } @@ -2474,10 +2489,16 @@ mod tests { async fn test_metadata_size_hint() { let store = Arc::new(object_store::memory::InMemory::new()) as Arc; - let store_url = ObjectStoreUrl::parse("memory://test").unwrap(); + let store_url = StorageUrl::parse("memory://test").unwrap(); let ctx = SessionContext::new(); - ctx.register_object_store(store_url.as_ref(), store.clone()); + ctx.register_storage( + store_url.as_ref(), + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + store.clone(), + )), + ) + .unwrap(); // write some data out, it doesn't matter what it is let c1: ArrayRef = Arc::new(Int32Array::from(vec![Some(1)])); @@ -2500,8 +2521,8 @@ mod tests { ); let config = FileScanConfigBuilder::new(store_url, source) .with_file( - PartitionedFile::new_from_meta(ObjectMeta { - location: Path::from(name_1), + PartitionedFile::new_from_meta(datafusion_storage::FileInfo { + location: datafusion_storage::path::Path::from(name_1), last_modified: Utc::now(), size: total_size_1, e_tag: None, @@ -2509,13 +2530,15 @@ mod tests { }) .with_metadata_size_hint(123), ) - .with_file(PartitionedFile::new_from_meta(ObjectMeta { - location: Path::from(name_2), - last_modified: Utc::now(), - size: total_size_2, - e_tag: None, - version: None, - })) + .with_file(PartitionedFile::new_from_meta( + datafusion_storage::FileInfo { + location: datafusion_storage::path::Path::from(name_2), + last_modified: Utc::now(), + size: total_size_2, + e_tag: None, + version: None, + }, + )) .build(); let exec = DataSourceExec::from_data_source(config); diff --git a/datafusion/core/src/execution/context/mod.rs b/datafusion/core/src/execution/context/mod.rs index ff1ad25811440..39ac66919090b 100644 --- a/datafusion/core/src/execution/context/mod.rs +++ b/datafusion/core/src/execution/context/mod.rs @@ -17,28 +17,29 @@ //! [`SessionContext`] API for registering data sources and executing queries +use crate::execution::SessionStateDefaults; use std::collections::HashSet; use std::fmt::Debug; use std::sync::{Arc, Weak}; use std::time::Duration; +use super::options::ArrowReadOptions; use super::options::ReadOptions; +use crate::catalog::listing_schema::ListingSchemaProvider; use crate::datasource::dynamic_file::DynamicListTableFactory; +use crate::datasource::listing::{ + ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl, +}; use crate::execution::session_state::SessionStateBuilder; use crate::{ - catalog::listing_schema::ListingSchemaProvider, catalog::{ CatalogProvider, CatalogProviderList, TableProvider, TableProviderFactory, }, dataframe::DataFrame, - datasource::listing::{ - ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl, - }, datasource::{MemTable, ViewTable, provider_as_source}, error::Result, execution::{ FunctionRegistry, - options::ArrowReadOptions, runtime_env::{RuntimeEnv, RuntimeEnvBuilder}, }, logical_expr::AggregateUDF, @@ -50,28 +51,28 @@ use crate::{ SetVariable, TableType, UNNAMED_TABLE, }, physical_expr::PhysicalExpr, - physical_plan::ExecutionPlan, variable::{VarProvider, VarType}, }; // backwards compatibility pub use crate::execution::session_state::SessionState; -use arrow::datatypes::{Schema, SchemaRef}; +use arrow::datatypes::Schema; +use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_catalog::MemoryCatalogProvider; use datafusion_catalog::memory::MemorySchemaProvider; -use datafusion_catalog::{ - DynamicFileCatalog, TableFunction, TableFunctionImpl, UrlTableFactory, -}; +use datafusion_catalog::{DynamicFileCatalog, UrlTableFactory}; +use datafusion_catalog::{TableFunction, TableFunctionImpl}; use datafusion_catalog_listing::SchemaSource; use datafusion_common::config::{ConfigField, ConfigOptions}; +use datafusion_common::internal_datafusion_err; use datafusion_common::metadata::ScalarAndMetadata; use datafusion_common::{ DFSchema, DataFusionError, ParamValues, SchemaError, SchemaReference, TableReference, config::{ConfigExtension, TableOptions}, - exec_datafusion_err, exec_err, internal_datafusion_err, not_impl_err, - plan_datafusion_err, plan_err, schema_err, + exec_datafusion_err, exec_err, not_impl_err, plan_datafusion_err, plan_err, + schema_err, tree_node::{TreeNodeRecursion, TreeNodeVisitor}, }; pub use datafusion_execution::TaskContext; @@ -101,9 +102,9 @@ use datafusion_optimizer::{Analyzer, OptimizerContext}; use datafusion_optimizer::{AnalyzerRule, OptimizerRule}; use datafusion_session::SessionStore; +use crate::physical_plan::ExecutionPlan; use async_trait::async_trait; use chrono::{DateTime, Utc}; -use object_store::ObjectStore; use parking_lot::RwLock; use url::Url; @@ -306,6 +307,24 @@ impl Default for SessionContext { } impl SessionContext { + /// Register a complete storage backend once for all users of this runtime. + /// Replacing a binding does not change already constructed storage tables. + pub fn register_storage( + &self, + url: &Url, + storage: Arc, + ) -> Result>> { + self.runtime_env().register_storage(url, storage) + } + + /// Deregister storage while allowing existing tables and plans to finish. + pub fn deregister_storage( + &self, + url: &Url, + ) -> Result> { + self.runtime_env().deregister_storage(url) + } + /// Creates a new `SessionContext` using the default [`SessionConfig`]. pub fn new() -> Self { Self::new_with_config(SessionConfig::new()) @@ -337,7 +356,7 @@ impl SessionContext { /// See [`Self::new_with_config_rt`] for more details on resource /// limits. pub fn new_with_config(config: SessionConfig) -> Self { - let runtime = Arc::new(RuntimeEnv::default()); + let runtime = SessionStateDefaults::default_runtime_env(); Self::new_with_config_rt(config, runtime) } @@ -502,36 +521,6 @@ impl SessionContext { self.state.write().add_analyzer_rule(analyzer_rule); } - /// Registers an [`ObjectStore`] to be used with a specific URL prefix. - /// - /// See [`RuntimeEnv::register_object_store`] for more details. - /// - /// # Example: register a local object store for the "file://" URL prefix - /// ``` - /// # use std::sync::Arc; - /// # use datafusion::prelude::SessionContext; - /// # use datafusion_execution::object_store::ObjectStoreUrl; - /// let object_store_url = ObjectStoreUrl::parse("file://").unwrap(); - /// let object_store = object_store::local::LocalFileSystem::new(); - /// let ctx = SessionContext::new(); - /// // All files with the file:// url prefix will be read from the local file system - /// ctx.register_object_store(object_store_url.as_ref(), Arc::new(object_store)); - /// ``` - pub fn register_object_store( - &self, - url: &Url, - object_store: Arc, - ) -> Option> { - self.runtime_env().register_object_store(url, object_store) - } - - /// Deregisters an [`ObjectStore`] associated with the specific URL prefix. - /// - /// See [`RuntimeEnv::deregister_object_store`] for more details. - pub fn deregister_object_store(&self, url: &Url) -> Result> { - self.runtime_env().deregister_object_store(url) - } - /// Registers the given [`RecordBatch`] as the specified table reference. pub fn register_batch( &self, @@ -1741,7 +1730,9 @@ impl SessionContext { } let schema_table_path = table_paths[0].clone(); + let storage = self.runtime_env().storage(&schema_table_path)?; let config = ListingTableConfig::new_with_multi_paths(table_paths) + .with_storage(storage) .with_listing_options(listing_options); let config = match options.schema_source() { SchemaSource::Inferred | SchemaSource::Unset => { @@ -1828,12 +1819,11 @@ impl SessionContext { )) } /// Registers a [`ListingTable`] that can assemble multiple files - /// from locations in an [`ObjectStore`] instance into a single + /// from locations in an [`StorageBinding`](datafusion_storage::StorageBinding) instance into a single /// table. /// /// This method is `async` because it might need to resolve the schema. /// - /// [`ObjectStore`]: object_store::ObjectStore pub async fn register_listing_table( &self, table_ref: impl Into, @@ -1845,13 +1835,14 @@ impl SessionContext { let table_ref = table_ref.into(); let table_path = ListingTableUrl::parse(table_path)?.with_table_ref(table_ref.clone()); - let resolved_schema = match provided_schema { - Some(s) => s, - None => options.infer_schema(&self.state(), &table_path).await?, - }; + let storage = self.runtime_env().storage(&table_path)?; let config = ListingTableConfig::new(table_path) - .with_listing_options(options) - .with_schema(resolved_schema); + .with_storage(storage) + .with_listing_options(options); + let config = match provided_schema { + Some(schema) => config.with_schema(schema), + None => config.infer_schema(&self.state()).await?, + }; let table = ListingTable::try_new(config)? .with_definition(sql_definition) .with_cache(self.runtime_env().cache_manager.get_file_statistic_cache()); @@ -2375,6 +2366,7 @@ mod tests { use crate::catalog::SchemaProvider; use crate::execution::session_state::SessionStateBuilder; + use crate::physical_plan::ExecutionPlan; use crate::physical_planner::PhysicalPlanner; use async_trait::async_trait; use datafusion_expr::planner::TypePlanner; @@ -2570,7 +2562,7 @@ mod tests { #[tokio::test] async fn custom_query_planner() -> Result<()> { - let runtime = Arc::new(RuntimeEnv::default()); + let runtime = SessionStateDefaults::default_runtime_env(); let session_state = SessionStateBuilder::new() .with_config(SessionConfig::new()) .with_runtime_env(runtime) diff --git a/datafusion/core/src/execution/session_state.rs b/datafusion/core/src/execution/session_state.rs index aa8ba4c3b733b..11bf864a581a6 100644 --- a/datafusion/core/src/execution/session_state.rs +++ b/datafusion/core/src/execution/session_state.rs @@ -45,8 +45,8 @@ use datafusion_common::config::{ConfigExtension, ConfigOptions, TableOptions}; use datafusion_common::display::{PlanType, StringifiedPlan, ToStringifiedPlan}; use datafusion_common::tree_node::TreeNode; use datafusion_common::{ - DFSchema, DataFusionError, ResolvedTableReference, TableReference, config_err, - exec_err, plan_datafusion_err, + DFSchema, DataFusionError, ResolvedTableReference, TableReference, exec_err, + plan_datafusion_err, }; use datafusion_execution::TaskContext; use datafusion_execution::config::SessionConfig; @@ -85,10 +85,11 @@ use datafusion_sql::{ use async_trait::async_trait; use chrono::{DateTime, Utc}; +use datafusion_common::config_err; use futures::future::BoxFuture; use itertools::Itertools; -use log::{debug, info}; -use object_store::ObjectStore; +use log::debug; +use log::info; #[cfg(feature = "sql")] use sqlparser::{ ast::{Expr as SQLExpr, ExprWithAlias as SQLExprWithAlias}, @@ -241,8 +242,9 @@ impl Debug for SessionState { .field("config", &self.inner.config) .field("runtime_env", &self.inner.runtime_env) .field("catalog_list", &self.inner.catalog_list) - .field("serializer_registry", &self.inner.serializer_registry) - .field("file_formats", &self.inner.file_formats) + .field("serializer_registry", &self.inner.serializer_registry); + let ret = ret.field("file_formats", &self.inner.file_formats); + let ret = ret .field("execution_props", &self.execution_props) .field("table_options", &self.inner.table_options) .field("table_factories", &self.inner.table_factories) @@ -1590,39 +1592,17 @@ impl SessionStateBuilder { self } - /// Register an `ObjectStore` to the [`RuntimeEnv`]. See [`RuntimeEnv::register_object_store`] - /// for more details. - /// - /// Note that this creates a default [`RuntimeEnv`] if there isn't one passed in already. - /// - /// ``` - /// # use datafusion::prelude::*; - /// # use datafusion::execution::session_state::SessionStateBuilder; - /// # use datafusion_execution::runtime_env::RuntimeEnv; - /// # use url::Url; - /// # use std::sync::Arc; - /// # let http_store = object_store::local::LocalFileSystem::new(); - /// let url = Url::try_from("file://").unwrap(); - /// let object_store = object_store::local::LocalFileSystem::new(); - /// let state = SessionStateBuilder::new() - /// .with_config(SessionConfig::new()) - /// .with_object_store(&url, Arc::new(object_store)) - /// .with_default_features() - /// .build(); - /// ``` - pub fn with_object_store( + /// Register a complete storage binding, creating the default runtime if needed. + pub fn with_storage( mut self, url: &Url, - object_store: Arc, - ) -> Self { - if self.runtime_env.is_none() { - self.runtime_env = Some(Arc::new(RuntimeEnv::default())); - } - self.runtime_env - .as_ref() - .unwrap() - .register_object_store(url, object_store); - self + storage: Arc, + ) -> datafusion_common::Result { + let runtime = self + .runtime_env + .get_or_insert_with(SessionStateDefaults::default_runtime_env); + runtime.register_storage(url, storage)?; + Ok(self) } /// Builds a [`SessionState`] with the current configuration. @@ -1665,7 +1645,9 @@ impl SessionStateBuilder { } = self; let config = config.unwrap_or_default(); - let runtime_env = runtime_env.unwrap_or_else(|| Arc::new(RuntimeEnv::default())); + let runtime_env = + runtime_env.unwrap_or_else(SessionStateDefaults::default_runtime_env); + SessionStateDefaults::configure_runtime(&runtime_env); let content = Arc::new(SessionStateInner { session_id: session_id.unwrap_or_else(|| Uuid::new_v4().to_string()), @@ -1984,8 +1966,9 @@ impl Debug for SessionStateBuilder { .field("config", &self.config) .field("runtime_env", &self.runtime_env) .field("catalog_list", &self.catalog_list) - .field("serializer_registry", &self.serializer_registry) - .field("file_formats", &self.file_formats) + .field("serializer_registry", &self.serializer_registry); + let ret = ret.field("file_formats", &self.file_formats); + let ret = ret .field("execution_props", &self.execution_props) .field("table_options", &self.table_options) .field("table_factories", &self.table_factories) diff --git a/datafusion/core/src/execution/session_state_defaults.rs b/datafusion/core/src/execution/session_state_defaults.rs index 59d3c440c7865..b78ac76a729f6 100644 --- a/datafusion/core/src/execution/session_state_defaults.rs +++ b/datafusion/core/src/execution/session_state_defaults.rs @@ -33,11 +33,11 @@ use crate::{functions, functions_aggregate, functions_table, functions_window}; use datafusion_catalog::TableFunction; use datafusion_catalog::{MemoryCatalogProvider, MemorySchemaProvider}; use datafusion_execution::config::SessionConfig; -use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_execution::runtime_env::RuntimeEnv; use datafusion_expr::planner::ExprPlanner; use datafusion_expr::registry::ExtensionTypeRegistrationRef; use datafusion_expr::{AggregateUDF, HigherOrderUDF, ScalarUDF, WindowUDF}; +use datafusion_storage::StorageUrl; use std::collections::HashMap; use std::sync::Arc; use url::Url; @@ -47,6 +47,26 @@ use url::Url; pub struct SessionStateDefaults {} impl SessionStateDefaults { + /// Assemble the default runtime. Backend features only affect this boundary. + pub fn default_runtime_env() -> Arc { + let runtime = Arc::new(RuntimeEnv::default()); + Self::configure_runtime(&runtime); + runtime + } + + /// Install default backends once, preserving application registrations. + pub(crate) fn configure_runtime(runtime: &RuntimeEnv) { + #[cfg(all(feature = "object_store", not(target_arch = "wasm32")))] + runtime + .storage_registry + .register_default( + StorageUrl::local_filesystem().as_ref(), + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::local()), + ) + .expect("valid local storage namespace"); + let _ = runtime; + } + /// returns a map of the default [`TableProviderFactory`]s pub fn default_table_factories() -> HashMap> { let mut table_factories: HashMap> = @@ -207,10 +227,10 @@ impl SessionStateDefaults { None => format!("{}://", url.scheme()), }; let path = &url.as_str()[authority.len()..]; - let path = object_store::path::Path::parse(path).expect("Can't parse path"); - let store = ObjectStoreUrl::parse(authority.as_str()) - .expect("Invalid default catalog url"); - let Ok(store) = runtime.object_store(store) else { + let path = datafusion_storage::path::Path::parse(path).expect("Can't parse path"); + let store = + StorageUrl::parse(authority.as_str()).expect("Invalid default catalog url"); + let Ok(store) = runtime.storage(store) else { return; }; let Some(factory) = table_factories.get(format.as_str()) else { diff --git a/datafusion/core/src/lib.rs b/datafusion/core/src/lib.rs index 28ea4b4490c3e..500c57c801db1 100644 --- a/datafusion/core/src/lib.rs +++ b/datafusion/core/src/lib.rs @@ -670,7 +670,7 @@ //! as the table definitions and the function registries. //! //! 2. [`TaskContext`]: State needed for execution such as the -//! [`MemoryPool`], [`DiskManager`], and [`ObjectStoreRegistry`]. +//! [`MemoryPool`], [`DiskManager`], and [`StorageRegistry`]. //! //! 3. [`ExecutionProps`]: Per-execution properties and data (such as //! starting timestamps, etc). @@ -689,7 +689,7 @@ //! [`DiskManager`]: crate::execution::DiskManager //! [`MemoryPool`]: crate::execution::memory_pool::MemoryPool //! [`RuntimeEnv`]: crate::execution::runtime_env::RuntimeEnv -//! [`ObjectStoreRegistry`]: crate::datasource::object_store::ObjectStoreRegistry +//! [`StorageRegistry`]: crate::storage::StorageRegistry //! //! ## Crate Organization //! @@ -779,7 +779,6 @@ pub mod scalar; // Re-export dependencies that are part of DataFusion public API (e.g. via DataFusionError) pub use arrow; -pub use object_store; #[cfg(feature = "parquet")] pub use parquet; @@ -1239,3 +1238,12 @@ doc_comment::doctest!( "../../../docs/source/contributor-guide/api-health.md", contributor_guide_api_health ); + +/// Backend-independent file access and storage registration. +pub mod storage { + pub use datafusion_storage::*; + #[cfg(feature = "object_store")] + pub use datafusion_storage_object_store::ObjectStoreStorage; + #[cfg(feature = "opendal")] + pub use datafusion_storage_opendal::OpendalStorage; +} diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 882956a63c114..9d1b086c7da98 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -682,7 +682,7 @@ impl DefaultPhysicalPlanner { let original_url = output_url.clone(); let input_exec = children.one()?; let parsed_url = ListingTableUrl::parse(output_url)?; - let object_store_url = parsed_url.object_store(); + let object_store_url = parsed_url.storage_url(); let schema = Arc::clone(input.schema().inner()); @@ -749,6 +749,9 @@ impl DefaultPhysicalPlanner { // Set file sink related options let config = FileSinkConfig { + storage: Some( + session_state.runtime_env().storage(&object_store_url)?, + ), original_url, object_store_url, table_paths: vec![parsed_url], diff --git a/datafusion/core/src/test/mod.rs b/datafusion/core/src/test/mod.rs index f46a5a0749065..3dc9c89817bd3 100644 --- a/datafusion/core/src/test/mod.rs +++ b/datafusion/core/src/test/mod.rs @@ -189,7 +189,7 @@ pub fn partitioned_file_groups( } } - // Must drop the stream before creating ObjectMeta below as drop triggers + // Must drop the stream before creating FileInfo below as drop triggers // finish for ZstdEncoder/BzEncoder which writes additional data for mut w in writers.into_iter() { w.flush().unwrap(); diff --git a/datafusion/core/src/test/object_store.rs b/datafusion/core/src/test/object_store.rs index 62c6699f8fcd1..3cb98c79dd30e 100644 --- a/datafusion/core/src/test/object_store.rs +++ b/datafusion/core/src/test/object_store.rs @@ -15,58 +15,49 @@ // specific language governing permissions and limitations // under the License. -//! Object store implementation used for testing - -use crate::{ - execution::{context::SessionState, session_state::SessionStateBuilder}, - object_store::{ - Error, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, - ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, - memory::InMemory, path::Path, - }, - prelude::SessionContext, -}; -use futures::{FutureExt, stream::BoxStream}; -use object_store::{CopyOptions, ObjectStoreExt}; -use std::{ - fmt::{Debug, Display, Formatter}, - sync::Arc, -}; -use tokio::{ - sync::Barrier, - time::{Duration, timeout}, +//! Native storage fixtures for listing tests. +use crate::execution::{context::SessionState, session_state::SessionStateBuilder}; +use crate::prelude::SessionContext; +use async_trait::async_trait; +use bytes::Bytes; +use datafusion_storage::{ + DirectoryListing, Error, FileAccessContext, FileInfo, FileReader, ReadRange, Storage, + StorageBinding, StorageUrl, path::Path, }; -use url::Url; +use futures::{StreamExt, stream::BoxStream}; +use std::{collections::BTreeMap, sync::Arc}; +use tokio::sync::Barrier; -/// Registers a test object store with the provided `ctx` +/// Register a read-only collection of zero-filled files. pub fn register_test_store(ctx: &SessionContext, files: &[(&str, u64)]) { - let url = Url::parse("test://").unwrap(); - ctx.register_object_store(&url, make_test_store_and_state(files).0); + let storage = make_test_store_and_state(files).0; + ctx.register_storage(storage.url().as_ref(), Arc::clone(storage.storage())) + .unwrap(); } -/// Create a test object store with the provided files -pub fn make_test_store_and_state(files: &[(&str, u64)]) -> (Arc, SessionState) { - let memory = InMemory::new(); - - for (name, size) in files { - memory - .put(&Path::from(*name), vec![0; *size as usize].into()) - .now_or_never() - .unwrap() - .unwrap(); - } - +/// Create a read-only collection and a session for listing tests. +pub fn make_test_store_and_state( + files: &[(&str, u64)], +) -> (Arc, SessionState) { + let files = files + .iter() + .map(|(path, size)| (Path::from(*path), *size)) + .collect(); + let backend = Arc::new(TestFiles { files }); ( - Arc::new(memory), + Arc::new(StorageBinding::new( + StorageUrl::parse("test://").unwrap(), + backend, + )), SessionStateBuilder::new().with_default_features().build(), ) } -/// Helper method to fetch the file size and date at given path and create a `ObjectMeta` -pub fn local_unpartitioned_file(path: impl AsRef) -> ObjectMeta { +/// Metadata of an existing local file. +pub fn local_unpartitioned_file(path: impl AsRef) -> FileInfo { let location = Path::from_filesystem_path(path.as_ref()).unwrap(); - let metadata = std::fs::metadata(path).expect("Local file metadata"); - ObjectMeta { + let metadata = std::fs::metadata(path).expect("local file metadata"); + FileInfo { location, last_modified: metadata.modified().map(chrono::DateTime::from).unwrap(), size: metadata.len(), @@ -74,119 +65,149 @@ pub fn local_unpartitioned_file(path: impl AsRef) -> ObjectMeta version: None, } } - -/// Blocks the object_store `head` call until `concurrency` number of calls are pending. -pub fn ensure_head_concurrency( - object_store: Arc, - concurrency: usize, -) -> Arc { - Arc::new(BlockingObjectStore::new(object_store, concurrency)) -} - -/// An object store that “blocks” in its `head` call until an expected number of concurrent calls are reached. #[derive(Debug)] -struct BlockingObjectStore { - inner: Arc, - barrier: Arc, +struct TestFiles { + files: BTreeMap, } - -impl BlockingObjectStore { - const NAME: &'static str = "BlockingObjectStore"; - fn new(inner: Arc, expected_concurrency: usize) -> Self { - Self { - inner, - barrier: Arc::new(Barrier::new(expected_concurrency)), - } +impl TestFiles { + fn descriptor(&self, path: &Path) -> datafusion_storage::Result { + let size = *self + .files + .get(path) + .ok_or_else(|| Error::NotFound(path.to_string()))?; + Ok(FileInfo::new(path.clone(), size)) } } - -impl Display for BlockingObjectStore { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - Display::fmt(&self.inner, f) +#[async_trait] +impl Storage for TestFiles { + async fn open( + &self, + path: &Path, + _: FileAccessContext, + ) -> datafusion_storage::Result> { + self.descriptor(path)?; + Ok(Arc::new(ZeroReader)) } -} -/// All trait methods are forwarded to the inner object store, except for -/// the `head` method which waits until the expected number of concurrent calls is reached. -#[async_trait::async_trait] -impl ObjectStore for BlockingObjectStore { - async fn put_opts( + async fn stat( &self, - location: &Path, - payload: PutPayload, - opts: PutOptions, - ) -> object_store::Result { - self.inner.put_opts(location, payload, opts).await + path: &Path, + _: &FileAccessContext, + ) -> datafusion_storage::Result { + self.descriptor(path) } - async fn put_multipart_opts( + fn list( &self, - location: &Path, - opts: PutMultipartOptions, - ) -> object_store::Result> { - self.inner.put_multipart_opts(location, opts).await + prefix: &Path, + _: FileAccessContext, + ) -> BoxStream<'_, datafusion_storage::Result> { + let files = self + .files + .keys() + .filter(|p| p.prefix_match(prefix).is_some()) + .map(|path| self.descriptor(path)) + .collect::>(); + futures::stream::iter(files).boxed() } - - async fn get_opts( + async fn list_with_delimiter( &self, - location: &Path, - options: GetOptions, - ) -> object_store::Result { - if options.head { - println!( - "{} received head call for {location}", - BlockingObjectStore::NAME - ); - // Wait until the expected number of concurrent calls is reached, but timeout after 1 second to avoid hanging failing tests. - let wait_result = timeout(Duration::from_secs(1), self.barrier.wait()).await; - match wait_result { - Ok(_) => println!( - "{} barrier reached for {location}", - BlockingObjectStore::NAME - ), - Err(_) => { - let error_message = format!( - "{} barrier wait timed out for {location}", - BlockingObjectStore::NAME - ); - log::error!("{error_message}"); - return Err(Error::Generic { - store: BlockingObjectStore::NAME, - source: error_message.into(), - }); - } + prefix: &Path, + _: &FileAccessContext, + ) -> datafusion_storage::Result { + let mut result = DirectoryListing::default(); + let mut directories = std::collections::BTreeSet::new(); + for path in self.files.keys() { + let Some(mut parts) = path.prefix_match(prefix) else { + continue; + }; + let Some(first) = parts.next() else { continue }; + if parts.next().is_some() { + directories.insert(prefix.clone().join(first)); + } else { + result.files.push(self.descriptor(path)?); } } - - // Forward the call to the inner object store. - self.inner.get_opts(location, options).await + result.directories = directories.into_iter().collect(); + Ok(result) } - fn delete_stream( - &self, - locations: BoxStream<'static, object_store::Result>, - ) -> BoxStream<'static, object_store::Result> { - self.inner.delete_stream(locations) +} +#[derive(Debug)] +struct ZeroReader; +#[async_trait] +impl FileReader for ZeroReader { + async fn read_range(&self, range: ReadRange) -> datafusion_storage::Result { + let length = match range { + ReadRange::Bounded(r) => r.end - r.start, + ReadRange::Suffix(n) => n, + }; + Ok(Bytes::from(vec![0; length as usize])) } - fn list( + async fn read_ranges( &self, - prefix: Option<&Path>, - ) -> BoxStream<'static, object_store::Result> { - self.inner.list(prefix) + ranges: Vec>, + ) -> datafusion_storage::Result> { + Ok(ranges + .into_iter() + .map(|r| Bytes::from(vec![0; (r.end - r.start) as usize])) + .collect()) + } + fn stream( + self: Arc, + range: Option, + ) -> BoxStream<'static, datafusion_storage::Result> { + futures::stream::once(async move { + self.read_range(range.unwrap_or(ReadRange::Bounded(0..0))) + .await + }) + .boxed() } +} - async fn list_with_delimiter( +/// Delay stat calls until the expected concurrency is reached. +pub fn ensure_head_concurrency( + inner: Arc, + concurrency: usize, +) -> Arc { + let url = inner.url().clone(); + let backend = Arc::new(BlockingDiscovery { + inner, + barrier: Barrier::new(concurrency), + }); + Arc::new(StorageBinding::new(url, backend)) +} +#[derive(Debug)] +struct BlockingDiscovery { + inner: Arc, + barrier: Barrier, +} +#[async_trait] +impl Storage for BlockingDiscovery { + async fn open( &self, - prefix: Option<&Path>, - ) -> object_store::Result { - self.inner.list_with_delimiter(prefix).await + path: &Path, + context: FileAccessContext, + ) -> datafusion_storage::Result> { + self.inner.storage().open(path, context).await } - async fn copy_opts( + async fn stat( + &self, + path: &Path, + context: &FileAccessContext, + ) -> datafusion_storage::Result { + tokio::time::timeout(std::time::Duration::from_secs(1), self.barrier.wait()) + .await + .map_err(|_| { + Error::InvalidInput("stat concurrency barrier timed out".into()) + })?; + self.inner.storage().stat(path, context).await + } + fn list( &self, - from: &Path, - to: &Path, - options: CopyOptions, - ) -> object_store::Result<()> { - self.inner.copy_opts(from, to, options).await + prefix: &Path, + context: FileAccessContext, + ) -> BoxStream<'_, datafusion_storage::Result> { + self.inner.storage().list(prefix, context) } } diff --git a/datafusion/core/src/test_util/mod.rs b/datafusion/core/src/test_util/mod.rs index ab7edef885b40..8c90cc2e16da1 100644 --- a/datafusion/core/src/test_util/mod.rs +++ b/datafusion/core/src/test_util/mod.rs @@ -40,17 +40,18 @@ use crate::error::Result; use crate::execution::session_state::CacheFactory; use crate::logical_expr::{LogicalPlanBuilder, UNNAMED_TABLE}; use crate::physical_plan::ExecutionPlan; -use crate::prelude::{CsvReadOptions, SessionContext}; +use crate::prelude::CsvReadOptions; +use crate::prelude::SessionContext; -use crate::execution::{SendableRecordBatchStream, SessionState, SessionStateBuilder}; +use crate::execution::SendableRecordBatchStream; +use crate::execution::{SessionState, SessionStateBuilder}; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use arrow::record_batch::RecordBatch; use datafusion_catalog::Session; -use datafusion_common::{DFSchemaRef, TableReference, plan_err}; -use datafusion_expr::{ - CreateExternalTable, Expr, LogicalPlan, SortExpr, TableType, - UserDefinedLogicalNodeCore, -}; +use datafusion_common::DFSchemaRef; +use datafusion_common::{TableReference, plan_err}; +use datafusion_expr::{CreateExternalTable, Expr, SortExpr, TableType}; +use datafusion_expr::{LogicalPlan, UserDefinedLogicalNodeCore}; use std::pin::Pin; use async_trait::async_trait; diff --git a/datafusion/core/src/test_util/parquet.rs b/datafusion/core/src/test_util/parquet.rs index b085010b0d73b..02b88d702d2df 100644 --- a/datafusion/core/src/test_util/parquet.rs +++ b/datafusion/core/src/test_util/parquet.rs @@ -25,7 +25,6 @@ use crate::arrow::{datatypes::SchemaRef, record_batch::RecordBatch}; use crate::common::ToDFSchema; use crate::config::ConfigOptions; use crate::datasource::listing::{ListingTableUrl, PartitionedFile}; -use crate::datasource::object_store::ObjectStoreUrl; use crate::datasource::physical_plan::ParquetSource; use crate::error::Result; use crate::logical_expr::execution_props::ExecutionProps; @@ -37,11 +36,12 @@ use crate::physical_plan::ExecutionPlan; use crate::physical_plan::filter::FilterExec; use crate::physical_plan::metrics::MetricsSet; use crate::prelude::{Expr, SessionConfig, SessionContext}; +use datafusion_storage::StorageUrl; use datafusion_datasource::file_scan_config::FileScanConfigBuilder; use datafusion_datasource::source::DataSourceExec; -use object_store::ObjectMeta; -use object_store::path::Path; +use datafusion_storage::FileInfo; +use datafusion_storage::path::Path; use parquet::arrow::ArrowWriter; use parquet::file::properties::WriterProperties; @@ -49,8 +49,8 @@ use parquet::file::properties::WriterProperties; pub struct TestParquetFile { path: PathBuf, schema: SchemaRef, - object_store_url: ObjectStoreUrl, - object_meta: ObjectMeta, + object_store_url: StorageUrl, + object_meta: FileInfo, } #[derive(Debug, Clone, Copy)] @@ -118,9 +118,9 @@ impl TestParquetFile { let object_store_url = ListingTableUrl::parse(canonical_path.to_str().unwrap_or_default())? - .object_store(); + .storage_url(); - let object_meta = ObjectMeta { + let object_meta = FileInfo { location: Path::parse(canonical_path.to_str().unwrap_or_default())?, last_modified: Default::default(), size, diff --git a/datafusion/core/tests/catalog_listing/pruned_partition_list.rs b/datafusion/core/tests/catalog_listing/pruned_partition_list.rs index 8f93dc17dbad2..f6188170187e3 100644 --- a/datafusion/core/tests/catalog_listing/pruned_partition_list.rs +++ b/datafusion/core/tests/catalog_listing/pruned_partition_list.rs @@ -19,7 +19,8 @@ use std::sync::Arc; use arrow_schema::DataType; use futures::{FutureExt, StreamExt as _, TryStreamExt as _}; -use object_store::{ObjectStoreExt, memory::InMemory, path::Path}; +use object_store::ObjectStoreExt; +use object_store::memory::InMemory; use datafusion::execution::SessionStateBuilder; use datafusion_catalog_listing::helpers::{ @@ -235,17 +236,23 @@ async fn test_list_partition() { pub fn make_test_store_and_state( files: &[(&str, u64)], -) -> (Arc, Arc) { +) -> (Arc, Arc) { let memory = InMemory::new(); for (name, size) in files { memory - .put(&Path::from(*name), vec![0; *size as usize].into()) + .put( + &object_store::path::Path::from(*name), + vec![0; *size as usize].into(), + ) .now_or_never() .unwrap() .unwrap(); } let state = SessionStateBuilder::new().build(); - (Arc::new(memory), Arc::new(state)) + ( + test_utils::storage::object_store(Arc::new(memory)), + Arc::new(state), + ) } diff --git a/datafusion/core/tests/dataframe/mod.rs b/datafusion/core/tests/dataframe/mod.rs index 9b0ca1b2768a8..88cdaaecdfde3 100644 --- a/datafusion/core/tests/dataframe/mod.rs +++ b/datafusion/core/tests/dataframe/mod.rs @@ -5889,9 +5889,15 @@ async fn write_partitioned_parquet_results(#[case] string_type: DataType) -> Res let ctx = SessionContext::new(); ctx.register_table("test", mem_table)?; - let local = Arc::new(LocalFileSystem::new_with_prefix(&tmp_dir)?); + let local = Arc::new(LocalFileSystem::new_with_prefix(&tmp_dir).unwrap()); let local_url = Url::parse("file://local").unwrap(); - ctx.register_object_store(&local_url, local); + ctx.register_storage( + &local_url, + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + local, + )), + ) + .unwrap(); // execute a simple query and write the results to parquet let out_dir = tmp_dir.as_ref().to_str().unwrap().to_string() + "/out/"; @@ -5970,9 +5976,15 @@ async fn write_parquet_results() -> Result<()> { .await?; // register a local file system object store for /tmp directory - let local = Arc::new(LocalFileSystem::new_with_prefix(&tmp_dir)?); + let local = Arc::new(LocalFileSystem::new_with_prefix(&tmp_dir).unwrap()); let local_url = Url::parse("file://local").unwrap(); - ctx.register_object_store(&local_url, local); + ctx.register_storage( + &local_url, + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + local, + )), + ) + .unwrap(); // execute a simple query and write the results to parquet let out_dir = tmp_dir.as_ref().to_str().unwrap().to_string() + "/out/"; diff --git a/datafusion/core/tests/datasource/object_store_access.rs b/datafusion/core/tests/datasource/object_store_access.rs index 16d894bde1303..875e3d14e9325 100644 --- a/datafusion/core/tests/datasource/object_store_access.rs +++ b/datafusion/core/tests/datasource/object_store_access.rs @@ -920,6 +920,98 @@ async fn query_single_parquet_file_multi_row_groups_multiple_predicates() { ); } +#[tokio::test] +async fn storage_adapter_preserves_native_read_requests() { + use datafusion_storage::{FileAccessContext, ReadRange}; + let test = Test::new().with_bytes("file", "0123456789").await; + let binding = test + .session_context + .runtime_env() + .storage_registry + .get(&Url::parse("mem://").unwrap()) + .unwrap(); + let reader = binding + .storage() + .open( + &datafusion_storage::path::Path::from("file"), + FileAccessContext::default(), + ) + .await + .unwrap(); + assert!(test.object_store.recorded_requests().is_empty()); + assert_eq!( + reader.read_range(ReadRange::Suffix(3)).await.unwrap(), + b"789"[..] + ); + assert_eq!(reader.read_range((1..3).into()).await.unwrap(), b"12"[..]); + reader.read_ranges(vec![0..2, 5..7]).await.unwrap(); + assert_eq!( + datafusion_storage::collect_bytes(reader.stream(None)) + .await + .unwrap(), + b"0123456789"[..] + ); + assert_snapshot!(test.requests(), @r" + RequestCountingObjectStore() + Total Requests: 4 + - GET (opts) path=file range=suffix:3 + - GET (opts) path=file range=1-3 + - GET (ranges) path=file ranges=0-2,5-7 + - GET (opts) path=file + "); + for request in test.object_store.recorded_requests() { + if let RequestDetails::GetOpts { get_options, .. } = request { + assert!(!get_options.head); + assert!(get_options.if_match.is_none()); + assert!(get_options.version.is_none()); + } + } +} + +#[tokio::test] +async fn storage_writer_retains_object_store_buffer_threshold() { + use datafusion_storage::{FileAccessContext, WriterOptions}; + use tokio::io::AsyncWriteExt; + for (size, multipart) in [(4, true), (1024, false)] { + let test = Test::new(); + let binding = test + .session_context + .runtime_env() + .storage_registry + .get(&Url::parse("mem://").unwrap()) + .unwrap(); + let mut writer = binding + .writer( + &datafusion_storage::path::Path::from("output"), + WriterOptions { + buffer_size: Some(size), + }, + FileAccessContext::default(), + ) + .await + .unwrap(); + writer.write_all(b"0123456789").await.unwrap(); + writer.shutdown().await.unwrap(); + let requests = test.object_store.recorded_requests(); + assert_eq!(requests.len(), 1); + assert_eq!( + matches!(requests[0], RequestDetails::PutMultipart), + multipart + ); + assert_eq!( + test.object_store + .inner + .get(&Path::from("output")) + .await + .unwrap() + .bytes() + .await + .unwrap(), + b"0123456789"[..] + ); + } +} + /// Runs tests with a request counting object store struct Test { object_store: Arc, @@ -942,7 +1034,13 @@ impl Test { let session_context = SessionContext::new(); session_context .runtime_env() - .register_object_store(&Url::parse("mem://").unwrap(), object_store.clone()); + .register_storage( + &Url::parse("mem://").unwrap(), + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + object_store.clone(), + )), + ) + .unwrap(); Self { object_store, session_context, @@ -1255,6 +1353,8 @@ impl Test { /// Details of individual requests made through the [`RequestCountingObjectStore`] #[derive(Clone, Debug)] enum RequestDetails { + Put, + PutMultipart, GetOpts { path: Path, get_options: GetOptions }, GetRanges { path: Path, ranges: Vec> }, List { prefix: Option }, @@ -1274,6 +1374,8 @@ fn display_range(range: &Range) -> impl Display + '_ { impl Display for RequestDetails { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { match self { + RequestDetails::Put => write!(f, "PUT"), + RequestDetails::PutMultipart => write!(f, "PUT (multipart)"), RequestDetails::GetOpts { path, get_options } => { write!(f, "GET (opts) path={path}")?; if let Some(range) = &get_options.range { @@ -1379,19 +1481,21 @@ impl RequestCountingObjectStore { impl ObjectStore for RequestCountingObjectStore { async fn put_opts( &self, - _location: &Path, - _payload: PutPayload, - _opts: PutOptions, + location: &Path, + payload: PutPayload, + opts: PutOptions, ) -> object_store::Result { - unimplemented!() + self.requests.lock().push(RequestDetails::Put); + self.inner.put_opts(location, payload, opts).await } async fn put_multipart_opts( &self, - _location: &Path, - _opts: PutMultipartOptions, + location: &Path, + opts: PutMultipartOptions, ) -> object_store::Result> { - unimplemented!() + self.requests.lock().push(RequestDetails::PutMultipart); + self.inner.put_multipart_opts(location, opts).await } async fn get_opts( diff --git a/datafusion/core/tests/fuzz_cases/pruning.rs b/datafusion/core/tests/fuzz_cases/pruning.rs index da7faffcb3a61..680bba1969f9d 100644 --- a/datafusion/core/tests/fuzz_cases/pruning.rs +++ b/datafusion/core/tests/fuzz_cases/pruning.rs @@ -27,9 +27,9 @@ use datafusion::{ use datafusion_common::DFSchema; use datafusion_datasource::file_scan_config::FileScanConfigBuilder; use datafusion_datasource::source::DataSourceExec; -use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_plan::{ExecutionPlan, collect, filter::FilterExec}; +use datafusion_storage::StorageUrl; use itertools::Itertools; use object_store::{ ObjectStore, ObjectStoreExt, PutPayload, memory::InMemory, path::Path, @@ -171,7 +171,13 @@ impl Utf8Test { } let store = Self::memory_store(); - ctx.register_object_store(&Url::parse("memory://").unwrap(), Arc::clone(store)); + ctx.register_storage( + &Url::parse("memory://").unwrap(), + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + Arc::clone(store), + )), + ) + .unwrap(); let files = Self::test_files().await; let schema = Self::schema(); @@ -278,7 +284,7 @@ async fn execute_with_predicate( ParquetSource::new(schema.clone()) }; let config = FileScanConfigBuilder::new( - ObjectStoreUrl::parse("memory://").unwrap(), + StorageUrl::parse("memory://").unwrap(), Arc::new(parquet_source), ) .with_file_group( diff --git a/datafusion/core/tests/fuzz_cases/topk_filter_pushdown.rs b/datafusion/core/tests/fuzz_cases/topk_filter_pushdown.rs index 80df91cb1036b..50c87fc4f7fbb 100644 --- a/datafusion/core/tests/fuzz_cases/topk_filter_pushdown.rs +++ b/datafusion/core/tests/fuzz_cases/topk_filter_pushdown.rs @@ -27,11 +27,12 @@ use datafusion::datasource::listing::{ListingOptions, ListingTable, ListingTable use datafusion::prelude::{SessionConfig, SessionContext}; use datafusion_datasource::ListingTableUrl; use datafusion_datasource_parquet::ParquetFormat; -use datafusion_execution::object_store::ObjectStoreUrl; +use datafusion_storage::StorageUrl; use itertools::Itertools; +use object_store::PutPayload; use object_store::memory::InMemory; use object_store::path::Path; -use object_store::{ObjectStore, ObjectStoreExt, PutPayload}; +use object_store::{ObjectStore, ObjectStoreExt}; use parquet::arrow::ArrowWriter; use rand::rngs::StdRng; use rand::{Rng, SeedableRng}; @@ -180,8 +181,14 @@ async fn run_query_with_config( let store = dataset.store; let schema = dataset.schema; let ctx = SessionContext::new_with_config(config); - let url = ObjectStoreUrl::parse("memory://").unwrap(); - ctx.register_object_store(url.as_ref(), store.clone()); + let url = StorageUrl::parse("memory://").unwrap(); + ctx.register_storage( + url.as_ref(), + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + store.clone(), + )), + ) + .unwrap(); let format = Arc::new( ParquetFormat::default() diff --git a/datafusion/core/tests/memory_limit/mod.rs b/datafusion/core/tests/memory_limit/mod.rs index 1636b676ba023..378d3a4f905ee 100644 --- a/datafusion/core/tests/memory_limit/mod.rs +++ b/datafusion/core/tests/memory_limit/mod.rs @@ -918,10 +918,10 @@ fn setup_context( .unwrap(); let runtime = Arc::new(RuntimeEnv { + storage_registry: Arc::clone(&runtime.storage_registry), memory_pool: runtime.memory_pool.clone(), disk_manager: Arc::new(disk_manager), cache_manager: runtime.cache_manager.clone(), - object_store_registry: runtime.object_store_registry.clone(), #[cfg(feature = "parquet_encryption")] parquet_encryption_factory_registry: runtime .parquet_encryption_factory_registry diff --git a/datafusion/core/tests/parquet/custom_reader.rs b/datafusion/core/tests/parquet/custom_reader.rs index 31337ed5d8485..c1727daf9e539 100644 --- a/datafusion/core/tests/parquet/custom_reader.rs +++ b/datafusion/core/tests/parquet/custom_reader.rs @@ -24,13 +24,13 @@ use arrow::array::{ArrayRef, Int8Array, Int64Array, StringArray}; use arrow::datatypes::{Field, Schema, SchemaBuilder}; use arrow::record_batch::RecordBatch; use datafusion::datasource::listing::PartitionedFile; -use datafusion::datasource::object_store::ObjectStoreUrl; use datafusion::datasource::physical_plan::{ ParquetFileMetrics, ParquetFileReaderFactory, ParquetSource, }; use datafusion::physical_plan::collect; use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; use datafusion::prelude::SessionContext; +use datafusion::storage::StorageUrl; use datafusion_common::Result; use datafusion_common::test_util::batches_to_sort_string; @@ -39,12 +39,13 @@ use datafusion_datasource::file_scan_config::FileScanConfigBuilder; use datafusion_datasource::source::DataSourceExec; use datafusion_datasource_parquet::metadata::DFParquetMetadata; use datafusion_datasource_parquet::{ParquetAccessPlan, RowGroupAccess}; +use datafusion_storage::path::Path; +use datafusion_storage::{FileAccessContext, FileInfo, StorageBinding}; +use futures::FutureExt; use futures::future::BoxFuture; -use futures::{FutureExt, TryFutureExt}; use insta::assert_snapshot; +use object_store::ObjectStoreExt; use object_store::memory::InMemory; -use object_store::path::Path; -use object_store::{ObjectMeta, ObjectStore, ObjectStoreExt}; use parquet::arrow::ArrowWriter; use parquet::arrow::arrow_reader::ArrowReaderOptions; use parquet::arrow::async_reader::AsyncFileReader; @@ -86,7 +87,7 @@ async fn route_data_access_ops_to_parquet_file_reader_factory() { ); let base_config = FileScanConfigBuilder::new( // just any url that doesn't point to in memory object store - ObjectStoreUrl::local_filesystem(), + StorageUrl::local_filesystem(), source, ) .with_file_group(file_group) @@ -137,10 +138,13 @@ async fn custom_payload_and_access_plan_coexist() { let location = Path::parse("two-row-groups.parquet").unwrap(); let size = buf.len() as u64; in_memory - .put(&location, Bytes::from(buf).into()) + .put( + &object_store::path::Path::from(location.as_ref()), + Bytes::from(buf).into(), + ) .await .unwrap(); - let meta = ObjectMeta { + let meta = FileInfo { location, last_modified: chrono::DateTime::from(SystemTime::now()), size, @@ -154,16 +158,15 @@ async fn custom_payload_and_access_plan_coexist() { .with_extension(String::from(EXPECTED_USER_DEFINED_METADATA)) .with_extension(access_plan); - let store: Arc = Arc::new(in_memory); + let store = test_utils::storage::object_store(Arc::new(in_memory)); let source = Arc::new( ParquetSource::new(file_schema.clone()).with_parquet_file_reader_factory( Arc::new(InMemoryParquetFileReaderFactory(Arc::clone(&store))), ), ); - let base_config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source) - .with_file_group(vec![pf].into()) - .build(); + let base_config = FileScanConfigBuilder::new(StorageUrl::local_filesystem(), source) + .with_file_group(vec![pf].into()) + .build(); let parquet_exec = DataSourceExec::from_data_source(base_config); let session_ctx = SessionContext::new(); @@ -186,7 +189,7 @@ async fn custom_payload_and_access_plan_coexist() { } #[derive(Debug)] -struct InMemoryParquetFileReaderFactory(Arc); +struct InMemoryParquetFileReaderFactory(Arc); impl ParquetFileReaderFactory for InMemoryParquetFileReaderFactory { fn create_reader( @@ -236,10 +239,10 @@ fn add_to_batch(batch: &RecordBatch, field_name: &str, array: ArrayRef) -> Recor async fn store_parquet_in_memory( batches: Vec, -) -> (Arc, Vec) { +) -> (Arc, Vec) { let in_memory = InMemory::new(); - let parquet_batches: Vec<(ObjectMeta, Bytes)> = batches + let parquet_batches: Vec<(FileInfo, Bytes)> = batches .into_iter() .enumerate() .map(|(offset, batch)| { @@ -252,7 +255,7 @@ async fn store_parquet_in_memory( writer.write(&batch).expect("Writing batch"); writer.close().unwrap(); - let meta = ObjectMeta { + let meta = FileInfo { location: Path::parse(format!("file-{offset}.parquet")) .expect("creating path"), last_modified: chrono::DateTime::from(SystemTime::now()), @@ -268,19 +271,25 @@ async fn store_parquet_in_memory( let mut objects = Vec::with_capacity(parquet_batches.len()); for (meta, bytes) in parquet_batches { in_memory - .put(&meta.location, bytes.into()) + .put( + &object_store::path::Path::from(meta.location.as_ref()), + bytes.into(), + ) .await .expect("put parquet file into in memory object store"); objects.push(meta); } - (Arc::new(in_memory), objects) + ( + test_utils::storage::object_store(Arc::new(in_memory)), + objects, + ) } /// Implements [`AsyncFileReader`] for a parquet file in object storage struct ParquetFileReader { - store: Arc, - meta: ObjectMeta, + store: Arc, + meta: FileInfo, metrics: ParquetFileMetrics, metadata_size_hint: Option, } @@ -293,12 +302,19 @@ impl AsyncFileReader for ParquetFileReader { let bytes_scanned = range.end - range.start; self.metrics.bytes_scanned.add(bytes_scanned as usize); - self.store - .get_range(&self.meta.location, range) - .map_err(|e| { - ParquetError::General(format!("AsyncChunkReader::get_bytes error: {e}")) - }) - .boxed() + async move { + let reader = self + .store + .open(&self.meta, FileAccessContext::new("custom-reader")) + .await + .map_err(|e| ParquetError::General(e.to_string()))?; + reader + .read_ranges(vec![range]) + .await + .map(|mut bytes| bytes.remove(0)) + .map_err(|e| ParquetError::General(e.to_string())) + } + .boxed() } fn get_metadata( @@ -306,7 +322,12 @@ impl AsyncFileReader for ParquetFileReader { _options: Option<&ArrowReaderOptions>, ) -> BoxFuture<'_, parquet::errors::Result>> { Box::pin(async move { - let metadata = DFParquetMetadata::new(self.store.as_ref(), &self.meta) + let reader = self + .store + .open(&self.meta, FileAccessContext::new("custom-reader")) + .await + .map_err(|e| ParquetError::General(e.to_string()))?; + let metadata = DFParquetMetadata::new(reader.as_ref(), &self.meta) .with_metadata_size_hint(self.metadata_size_hint) .fetch_metadata() .await diff --git a/datafusion/core/tests/parquet/encryption.rs b/datafusion/core/tests/parquet/encryption.rs index b7bfcefa30f34..8b2a501a89e50 100644 --- a/datafusion/core/tests/parquet/encryption.rs +++ b/datafusion/core/tests/parquet/encryption.rs @@ -298,7 +298,7 @@ async fn verify_file_encrypted( file_path.to_str().unwrap().to_owned() }; - let object_path = object_store::path::Path::from(file_path_str); + let object_path = datafusion::storage::path::Path::from(file_path_str); let decryption_properties = encryption_factory .get_file_decryption_properties(&options, &object_path) .await? @@ -326,7 +326,7 @@ async fn verify_file_encrypted( /// which generates encryption keys in a sequence #[derive(Debug, Default)] struct MockEncryptionFactory { - pub encryption_keys: Mutex>>, + pub encryption_keys: Mutex>>, pub counter: AtomicU8, } @@ -336,7 +336,7 @@ impl EncryptionFactory for MockEncryptionFactory { &self, config: &EncryptionFactoryOptions, _schema: &SchemaRef, - file_path: &object_store::path::Path, + file_path: &datafusion::storage::path::Path, ) -> datafusion_common::Result>> { assert_eq!( config.options.get("test_key"), @@ -353,7 +353,7 @@ impl EncryptionFactory for MockEncryptionFactory { async fn get_file_decryption_properties( &self, config: &EncryptionFactoryOptions, - file_path: &object_store::path::Path, + file_path: &datafusion::storage::path::Path, ) -> datafusion_common::Result>> { assert_eq!( config.options.get("test_key"), diff --git a/datafusion/core/tests/parquet/expr_adapter.rs b/datafusion/core/tests/parquet/expr_adapter.rs index dfed2eb5bda74..15ab527c712b8 100644 --- a/datafusion/core/tests/parquet/expr_adapter.rs +++ b/datafusion/core/tests/parquet/expr_adapter.rs @@ -36,13 +36,13 @@ use datafusion_common::DataFusionError; use datafusion_common::ScalarValue; use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; use datafusion_datasource::ListingTableUrl; -use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_expr::expressions::{self, Column}; use datafusion_physical_expr_adapter::{ DefaultPhysicalExprAdapter, DefaultPhysicalExprAdapterFactory, PhysicalExprAdapter, PhysicalExprAdapterFactory, }; +use datafusion_storage::StorageUrl; use object_store::{ObjectStore, ObjectStoreExt, memory::InMemory, path::Path}; use parquet::arrow::ArrowWriter; @@ -258,8 +258,14 @@ async fn register_memory_listing_table( base_path: &str, table_schema: SchemaRef, ) { - let store_url = ObjectStoreUrl::parse("memory://").unwrap(); - ctx.register_object_store(store_url.as_ref(), Arc::clone(&store)); + let store_url = StorageUrl::parse("memory://").unwrap(); + ctx.register_storage( + store_url.as_ref(), + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + store.clone(), + )), + ) + .unwrap(); let listing_table_config = ListingTableConfig::new(ListingTableUrl::parse(base_path).unwrap()) @@ -601,7 +607,7 @@ async fn test_custom_schema_adapter_and_custom_expression_adapter() { record_batch!(("extra", Int64, [1, 2, 3]), ("c1", Int32, [1, 2, 3])).unwrap(); let store = Arc::new(InMemory::new()) as Arc; - let store_url = ObjectStoreUrl::parse("memory://").unwrap(); + let store_url = StorageUrl::parse("memory://").unwrap(); let path = "test.parquet"; write_parquet(batch, store.clone(), path).await; @@ -617,7 +623,13 @@ async fn test_custom_schema_adapter_and_custom_expression_adapter() { .with_parquet_page_index_pruning(false); cfg.options_mut().execution.parquet.pushdown_filters = true; let ctx = SessionContext::new_with_config(cfg); - ctx.register_object_store(store_url.as_ref(), Arc::clone(&store)); + ctx.register_storage( + store_url.as_ref(), + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + store.clone(), + )), + ) + .unwrap(); assert!( !ctx.state() .config_mut() @@ -699,7 +711,7 @@ async fn test_physical_expr_adapter_with_non_null_defaults() { let batch = record_batch!(("c1", Int32, [10, 20, 30])).unwrap(); let store = Arc::new(InMemory::new()) as Arc; - let store_url = ObjectStoreUrl::parse("memory://").unwrap(); + let store_url = StorageUrl::parse("memory://").unwrap(); write_parquet(batch, store.clone(), "defaults_test.parquet").await; // Table schema has additional columns c2 (Utf8) and c3 (Int64) that don't exist in file @@ -714,7 +726,13 @@ async fn test_physical_expr_adapter_with_non_null_defaults() { .with_parquet_pruning(false); cfg.options_mut().execution.parquet.pushdown_filters = true; let ctx = SessionContext::new_with_config(cfg); - ctx.register_object_store(store_url.as_ref(), Arc::clone(&store)); + ctx.register_storage( + store_url.as_ref(), + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + store.clone(), + )), + ) + .unwrap(); // CustomPhysicalExprAdapterFactory fills: // - missing Utf8 columns with 'b' @@ -1026,7 +1044,7 @@ async fn test_struct_schema_evolution_projection_and_filter() -> Result<()> { RecordBatch::try_new(Arc::clone(&physical_schema), vec![Arc::new(struct_array)])?; let store = Arc::new(InMemory::new()) as Arc; - let store_url = ObjectStoreUrl::parse("memory://").unwrap(); + let store_url = StorageUrl::parse("memory://").unwrap(); write_parquet(batch, store.clone(), "struct_evolution.parquet").await; // Logical struct: {id: Int64?, name: Utf8?, extra: Boolean?} + metadata @@ -1051,7 +1069,13 @@ async fn test_struct_schema_evolution_projection_and_filter() -> Result<()> { cfg.options_mut().execution.parquet.pushdown_filters = true; let ctx = SessionContext::new_with_config(cfg); - ctx.register_object_store(store_url.as_ref(), Arc::clone(&store)); + ctx.register_storage( + store_url.as_ref(), + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + store.clone(), + )), + ) + .unwrap(); let listing_table_config = ListingTableConfig::new(ListingTableUrl::parse("memory:///").unwrap()) @@ -1452,7 +1476,7 @@ async fn test_physical_expr_adapter_factory_reuse_across_tables() { let batch2 = record_batch!(("c1", Int32, [10, 20, 30])).unwrap(); let store = Arc::new(InMemory::new()) as Arc; - let store_url = ObjectStoreUrl::parse("memory://").unwrap(); + let store_url = StorageUrl::parse("memory://").unwrap(); // Write files to different paths write_parquet(batch1, store.clone(), "table1/data.parquet").await; @@ -1469,7 +1493,13 @@ async fn test_physical_expr_adapter_factory_reuse_across_tables() { .with_parquet_pruning(false); cfg.options_mut().execution.parquet.pushdown_filters = true; let ctx = SessionContext::new_with_config(cfg); - ctx.register_object_store(store_url.as_ref(), Arc::clone(&store)); + ctx.register_storage( + store_url.as_ref(), + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + store.clone(), + )), + ) + .unwrap(); // Create ONE factory instance wrapped in Arc - this will be REUSED let factory: Arc = diff --git a/datafusion/core/tests/parquet/external_access_plan.rs b/datafusion/core/tests/parquet/external_access_plan.rs index 4d74a1f28264a..82a3d1e891af2 100644 --- a/datafusion/core/tests/parquet/external_access_plan.rs +++ b/datafusion/core/tests/parquet/external_access_plan.rs @@ -33,10 +33,10 @@ use datafusion_common::{DFSchema, assert_batches_eq, assert_contains}; use datafusion_datasource_parquet::{ ParquetAccessPlan, ParquetRowSelection, RowGroupAccess, }; -use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_expr::{Expr, col, lit}; use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::metrics::{MetricValue, MetricsSet}; +use datafusion_storage::StorageUrl; use datafusion_datasource::file_scan_config::FileScanConfigBuilder; use datafusion_datasource::source::DataSourceExec; @@ -460,7 +460,7 @@ impl TestFull { } // Create a DataSourceExec to read the file - let object_store_url = ObjectStoreUrl::local_filesystem(); + let object_store_url = StorageUrl::local_filesystem(); // add the predicate, if requested let source = if let Some(predicate) = predicate { let df_schema = DFSchema::try_from(schema.clone())?; diff --git a/datafusion/core/tests/parquet/page_pruning.rs b/datafusion/core/tests/parquet/page_pruning.rs index 372a7a601d492..345bfabc795cc 100644 --- a/datafusion/core/tests/parquet/page_pruning.rs +++ b/datafusion/core/tests/parquet/page_pruning.rs @@ -25,13 +25,13 @@ use arrow::datatypes::{DataType, Field, Schema}; use datafusion::datasource::file_format::FileFormat; use datafusion::datasource::file_format::parquet::ParquetFormat; use datafusion::datasource::listing::PartitionedFile; -use datafusion::datasource::object_store::ObjectStoreUrl; use datafusion::datasource::physical_plan::ParquetSource; use datafusion::datasource::source::DataSourceExec; use datafusion::execution::context::SessionState; use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::metrics::MetricValue; use datafusion::prelude::{SessionConfig, SessionContext}; +use datafusion::storage::StorageUrl; use datafusion_common::{ScalarValue, ToDFSchema}; use datafusion_expr::execution_props::ExecutionProps; use datafusion_expr::{Expr, col, lit}; @@ -39,9 +39,9 @@ use datafusion_physical_expr::create_physical_expr; use datafusion_datasource::file_scan_config::FileScanConfigBuilder; use datafusion_expr::physical_planning_context::PhysicalPlanningContext; +use datafusion_storage::FileInfo; +use datafusion_storage::path::Path; use futures::StreamExt; -use object_store::ObjectMeta; -use object_store::path::Path; use parquet::arrow::ArrowWriter; use parquet::file::properties::WriterProperties; @@ -50,15 +50,15 @@ async fn get_parquet_exec( filter: Expr, pushdown_filters: bool, ) -> DataSourceExec { - let object_store_url = ObjectStoreUrl::local_filesystem(); - let store = state.runtime_env().object_store(&object_store_url).unwrap(); + let object_store_url = StorageUrl::local_filesystem(); + let store = state.runtime_env().storage(&object_store_url).unwrap(); let testdata = datafusion::test_util::parquet_test_data(); let filename = format!("{testdata}/alltypes_tiny_pages.parquet"); let location = Path::from_filesystem_path(filename.as_str()).unwrap(); let metadata = std::fs::metadata(filename).expect("Local file metadata"); - let meta = ObjectMeta { + let meta = FileInfo { location, last_modified: metadata.modified().map(chrono::DateTime::from).unwrap(), size: metadata.len(), diff --git a/datafusion/core/tests/parquet/schema_coercion.rs b/datafusion/core/tests/parquet/schema_coercion.rs index be45ab38dabad..702d9bf88b0ea 100644 --- a/datafusion/core/tests/parquet/schema_coercion.rs +++ b/datafusion/core/tests/parquet/schema_coercion.rs @@ -28,12 +28,12 @@ use datafusion::prelude::SessionContext; use datafusion::test::object_store::local_unpartitioned_file; use datafusion_common::Result; use datafusion_common::test_util::batches_to_sort_string; -use datafusion_execution::object_store::ObjectStoreUrl; +use datafusion_storage::StorageUrl; use datafusion_datasource::file_scan_config::FileScanConfigBuilder; use datafusion_datasource::source::DataSourceExec; +use datafusion_storage::FileInfo; use insta::assert_snapshot; -use object_store::ObjectMeta; use parquet::arrow::ArrowWriter; use parquet::file::properties::WriterProperties; use tempfile::NamedTempFile; @@ -63,7 +63,7 @@ async fn multi_parquet_coercion() { Field::new("c3", DataType::Float64, true), ])); let source = Arc::new(ParquetSource::new(file_schema.clone())); - let conf = FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source) + let conf = FileScanConfigBuilder::new(StorageUrl::local_filesystem(), source) .with_file_group(file_group) .build(); @@ -117,7 +117,7 @@ async fn multi_parquet_coercion_projection() { Field::new("c3", DataType::Float64, true), ])); let config = FileScanConfigBuilder::new( - ObjectStoreUrl::local_filesystem(), + StorageUrl::local_filesystem(), Arc::new(ParquetSource::new(file_schema)), ) .with_file_group(file_group) @@ -148,7 +148,7 @@ async fn multi_parquet_coercion_projection() { /// Writes `batches` to a temporary parquet file pub fn store_parquet( batches: Vec, -) -> Result<(Vec, Vec)> { +) -> Result<(Vec, Vec)> { // Each batch writes to their own file let files: Vec<_> = batches .into_iter() diff --git a/datafusion/core/tests/parquet/string_in_list_pruning.rs b/datafusion/core/tests/parquet/string_in_list_pruning.rs index f2646855713ce..8a45b08363d16 100644 --- a/datafusion/core/tests/parquet/string_in_list_pruning.rs +++ b/datafusion/core/tests/parquet/string_in_list_pruning.rs @@ -29,18 +29,18 @@ use arrow::datatypes::{DataType, Field, Schema}; use arrow::record_batch::RecordBatch; use arrow::util::pretty::pretty_format_batches; use datafusion::datasource::listing::PartitionedFile; -use datafusion::datasource::object_store::ObjectStoreUrl; use datafusion::datasource::physical_plan::ParquetSource; use datafusion::datasource::source::DataSourceExec; use datafusion::physical_plan::{ExecutionPlan, collect, displayable}; use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext}; +use datafusion::storage::StorageUrl; use datafusion_common::config::TableParquetOptions; use datafusion_common::{ScalarValue, assert_batches_eq}; use datafusion_datasource::file_scan_config::FileScanConfigBuilder; use datafusion_physical_expr::expressions::{col, in_list, lit}; use datafusion_physical_plan::filter::FilterExec; use datafusion_physical_plan::metrics::{MetricValue, MetricsSet}; -use object_store::path::Path; +use datafusion_storage::path::Path; use parquet::arrow::ArrowWriter; use parquet::file::metadata::ParquetMetaData; use parquet::file::properties::{EnabledStatistics, WriterProperties}; @@ -340,7 +340,7 @@ async fn scan_ordered( location.to_string(), file.as_file().metadata().unwrap().len(), ); - let config = FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source) + let config = FileScanConfigBuilder::new(StorageUrl::local_filesystem(), source) .with_file(partitioned_file) .build(); let scan: Arc = Arc::new(DataSourceExec::new(Arc::new(config))); @@ -709,7 +709,7 @@ async fn check_in_list_with_null_preserves_filter_semantics( .with_bloom_filter_on_read(false), ); let config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source) + FileScanConfigBuilder::new(StorageUrl::local_filesystem(), source) .with_file(partitioned_file.clone()) .with_limit(Some(1)) .build(); diff --git a/datafusion/core/tests/physical_optimizer/aggregate_statistics.rs b/datafusion/core/tests/physical_optimizer/aggregate_statistics.rs index 2d22b60856ca5..60ee8616aca8c 100644 --- a/datafusion/core/tests/physical_optimizer/aggregate_statistics.rs +++ b/datafusion/core/tests/physical_optimizer/aggregate_statistics.rs @@ -36,7 +36,6 @@ use datafusion_common::{ColumnStatistics, Result, Statistics}; use datafusion_common::{ScalarValue, assert_batches_eq}; use datafusion_datasource::file_scan_config::FileScanConfigBuilder; use datafusion_execution::TaskContext; -use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_expr::Operator; use datafusion_functions_aggregate::count::count_udaf; use datafusion_functions_aggregate::sum::sum_udaf; @@ -53,6 +52,7 @@ use datafusion_physical_plan::common; use datafusion_physical_plan::displayable; use datafusion_physical_plan::filter::FilterExec; use datafusion_physical_plan::projection::ProjectionExec; +use datafusion_storage::StorageUrl; /// Mock data using a MemorySourceConfig which has an exact count statistic fn mock_data() -> Result> { @@ -472,7 +472,7 @@ async fn test_count_distinct_optimization() -> Result<()> { }; let config = FileScanConfigBuilder::new( - ObjectStoreUrl::parse("test:///").unwrap(), + StorageUrl::parse("test:///").unwrap(), Arc::new(ParquetSource::new(Arc::clone(&schema))), ) .with_file(PartitionedFile::new("x".to_string(), 100)) @@ -762,7 +762,7 @@ async fn test_sum_from_statistics() -> Result<()> { }; let config = FileScanConfigBuilder::new( - ObjectStoreUrl::parse("test:///").unwrap(), + StorageUrl::parse("test:///").unwrap(), Arc::new(ParquetSource::new(Arc::clone(&schema))), ) .with_file(PartitionedFile::new("x".to_string(), 100)) diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 6b8e38d150c23..0486f2a6cb36a 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -33,10 +33,10 @@ use datafusion::config::ConfigOptions; use datafusion::datasource::MemTable; use datafusion::datasource::file_format::file_compression_type::FileCompressionType; use datafusion::datasource::listing::PartitionedFile; -use datafusion::datasource::object_store::ObjectStoreUrl; use datafusion::datasource::physical_plan::{CsvSource, ParquetSource}; use datafusion::datasource::source::DataSourceExec; use datafusion::prelude::{SessionConfig, SessionContext}; +use datafusion::storage::StorageUrl; use datafusion_common::ScalarValue; use datafusion_common::Statistics; use datafusion_common::config::CsvOptions; @@ -364,7 +364,7 @@ fn parquet_exec_multiple_sorted( output_ordering: Vec, ) -> Arc { let config = FileScanConfigBuilder::new( - ObjectStoreUrl::parse("test:///").unwrap(), + StorageUrl::parse("test:///").unwrap(), Arc::new(ParquetSource::new(schema())), ) .with_file_groups(vec![ @@ -387,7 +387,7 @@ fn parquet_exec_with_output_partitioning( .collect::>(); let config = FileScanConfigBuilder::new( - ObjectStoreUrl::parse("test:///").unwrap(), + StorageUrl::parse("test:///").unwrap(), Arc::new(ParquetSource::new(schema())), ) .with_file_groups(file_groups) @@ -423,19 +423,18 @@ fn csv_exec() -> Arc { } fn csv_exec_with_sort(output_ordering: Vec) -> Arc { - let config = - FileScanConfigBuilder::new(ObjectStoreUrl::parse("test:///").unwrap(), { - let options = CsvOptions { - has_header: Some(false), - delimiter: b',', - quote: b'"', - ..Default::default() - }; - Arc::new(CsvSource::new(schema()).with_csv_options(options)) - }) - .with_file(PartitionedFile::new("x".to_string(), 100)) - .with_output_ordering(output_ordering) - .build(); + let config = FileScanConfigBuilder::new(StorageUrl::parse("test:///").unwrap(), { + let options = CsvOptions { + has_header: Some(false), + delimiter: b',', + quote: b'"', + ..Default::default() + }; + Arc::new(CsvSource::new(schema()).with_csv_options(options)) + }) + .with_file(PartitionedFile::new("x".to_string(), 100)) + .with_output_ordering(output_ordering) + .build(); DataSourceExec::from_data_source(config) } @@ -446,22 +445,21 @@ fn csv_exec_multiple() -> Arc { // Created a sorted parquet exec with multiple files fn csv_exec_multiple_sorted(output_ordering: Vec) -> Arc { - let config = - FileScanConfigBuilder::new(ObjectStoreUrl::parse("test:///").unwrap(), { - let options = CsvOptions { - has_header: Some(false), - delimiter: b',', - quote: b'"', - ..Default::default() - }; - Arc::new(CsvSource::new(schema()).with_csv_options(options)) - }) - .with_file_groups(vec![ - FileGroup::new(vec![PartitionedFile::new("x".to_string(), 100)]), - FileGroup::new(vec![PartitionedFile::new("y".to_string(), 100)]), - ]) - .with_output_ordering(output_ordering) - .build(); + let config = FileScanConfigBuilder::new(StorageUrl::parse("test:///").unwrap(), { + let options = CsvOptions { + has_header: Some(false), + delimiter: b',', + quote: b'"', + ..Default::default() + }; + Arc::new(CsvSource::new(schema()).with_csv_options(options)) + }) + .with_file_groups(vec![ + FileGroup::new(vec![PartitionedFile::new("x".to_string(), 100)]), + FileGroup::new(vec![PartitionedFile::new("y".to_string(), 100)]), + ]) + .with_output_ordering(output_ordering) + .build(); DataSourceExec::from_data_source(config) } @@ -2928,7 +2926,7 @@ fn parquet_exec_with_size( statistics.num_rows = Precision::Inexact(num_rows); statistics.total_byte_size = Precision::Inexact(total_byte_size); let config = FileScanConfigBuilder::new( - ObjectStoreUrl::parse("test:///").unwrap(), + StorageUrl::parse("test:///").unwrap(), Arc::new(ParquetSource::new(schema())), ) .with_file(PartitionedFile::new( @@ -3670,7 +3668,7 @@ fn parallelization_compressed_csv() -> Result<()> { for compression_type in compression_types { let plan = aggregate_exec_with_alias( DataSourceExec::from_data_source( - FileScanConfigBuilder::new(ObjectStoreUrl::parse("test:///").unwrap(), { + FileScanConfigBuilder::new(StorageUrl::parse("test:///").unwrap(), { let options = CsvOptions { has_header: Some(false), delimiter: b',', diff --git a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs index 04019feb6689a..2c095000841c7 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs @@ -38,7 +38,7 @@ use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuil use datafusion_datasource::source::DataSourceExec; use datafusion_expr_common::operator::Operator; use datafusion_expr::{JoinType, SortExpr}; -use datafusion_execution::object_store::ObjectStoreUrl; +use datafusion_storage::StorageUrl; use datafusion_physical_expr_common::sort_expr::{ LexOrdering, PhysicalSortExpr, PhysicalSortRequirement, OrderingRequirements }; @@ -86,7 +86,7 @@ fn csv_exec_sorted( ..Default::default() }; let mut builder = FileScanConfigBuilder::new( - ObjectStoreUrl::parse("test:///").unwrap(), + StorageUrl::parse("test:///").unwrap(), Arc::new(CsvSource::new(schema.clone()).with_csv_options(options)), ) .with_file(PartitionedFile::new("x".to_string(), 100)); diff --git a/datafusion/core/tests/physical_optimizer/filter_pushdown.rs b/datafusion/core/tests/physical_optimizer/filter_pushdown.rs index 26e6e0c74c49d..166e720aab007 100644 --- a/datafusion/core/tests/physical_optimizer/filter_pushdown.rs +++ b/datafusion/core/tests/physical_optimizer/filter_pushdown.rs @@ -42,7 +42,6 @@ use datafusion_common::{ use datafusion_datasource::{ PartitionedFile, file_groups::FileGroup, file_scan_config::FileScanConfigBuilder, }; -use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_expr::ScalarUDF; use datafusion_functions::math::random::RandomFunc; use datafusion_functions_aggregate::{ @@ -75,6 +74,7 @@ use datafusion_physical_plan::{ repartition::RepartitionExec, sorts::sort::SortExec, }; +use datafusion_storage::StorageUrl; use super::pushdown_utils::{ OptimizationTest, TestNode, TestScanBuilder, TestSource, format_plan_for_test, @@ -900,7 +900,7 @@ async fn test_topk_filter_passes_through_coalesce_partitions() { let source = Arc::new(TestSource::new(schema(), true, batches)); let base_config = - FileScanConfigBuilder::new(ObjectStoreUrl::parse("test://").unwrap(), source) + FileScanConfigBuilder::new(StorageUrl::parse("test://").unwrap(), source) .with_file_groups(vec![ // Partition 0 FileGroup::new(vec![PartitionedFile::new("test1.parquet", 123)]), @@ -998,10 +998,14 @@ async fn optimize_and_collect_pushdown_plan( .unwrap(); let session_ctx = SessionContext::new_with_config(SessionConfig::from(config).with_batch_size(10)); - session_ctx.register_object_store( - ObjectStoreUrl::parse("test://").unwrap().as_ref(), - Arc::new(InMemory::new()), - ); + session_ctx + .register_storage( + StorageUrl::parse("test://").unwrap().as_ref(), + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + Arc::new(InMemory::new()), + )), + ) + .unwrap(); let task_ctx = session_ctx.state().task_ctx(); let batches = collect(Arc::clone(&plan), task_ctx).await.unwrap(); (plan, batches) @@ -2287,10 +2291,14 @@ async fn test_no_pushdown_through_global_aggregate_with_name_collision() { assert!(optimized.downcast_ref::().is_some()); let session_ctx = SessionContext::new(); - session_ctx.register_object_store( - ObjectStoreUrl::parse("test://").unwrap().as_ref(), - Arc::new(InMemory::new()), - ); + session_ctx + .register_storage( + StorageUrl::parse("test://").unwrap().as_ref(), + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + Arc::new(InMemory::new()), + )), + ) + .unwrap(); let batches = collect(optimized, session_ctx.state().task_ctx()) .await .unwrap(); @@ -2728,10 +2736,14 @@ async fn test_hashjoin_dynamic_filter_all_partitions_empty() { // Put some data through the plan to check that the filter is updated to reflect the TopK state let session_ctx = SessionContext::new_with_config(config); - session_ctx.register_object_store( - ObjectStoreUrl::parse("test://").unwrap().as_ref(), - Arc::new(InMemory::new()), - ); + session_ctx + .register_storage( + StorageUrl::parse("test://").unwrap().as_ref(), + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + Arc::new(InMemory::new()), + )), + ) + .unwrap(); let state = session_ctx.state(); let task_ctx = state.task_ctx(); // Execute all partitions (required for partitioned hash join coordination) @@ -3128,10 +3140,14 @@ async fn test_hashjoin_dynamic_filter_survives_probe_subtree_replacement() { let session_ctx = SessionContext::new_with_config(SessionConfig::from(config).with_batch_size(10)); - session_ctx.register_object_store( - ObjectStoreUrl::parse("test://").unwrap().as_ref(), - Arc::new(InMemory::new()), - ); + session_ctx + .register_storage( + StorageUrl::parse("test://").unwrap().as_ref(), + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + Arc::new(InMemory::new()), + )), + ) + .unwrap(); collect(Arc::clone(&plan), session_ctx.state().task_ctx()) .await .unwrap(); diff --git a/datafusion/core/tests/physical_optimizer/projection_pushdown.rs b/datafusion/core/tests/physical_optimizer/projection_pushdown.rs index 10bdcd2d5d6b0..5d4d38690c2f2 100644 --- a/datafusion/core/tests/physical_optimizer/projection_pushdown.rs +++ b/datafusion/core/tests/physical_optimizer/projection_pushdown.rs @@ -27,7 +27,6 @@ use datafusion_common::config::{ConfigOptions, CsvOptions}; use datafusion_common::{JoinSide, JoinType, NullEquality, Result, ScalarValue}; use datafusion_datasource::TableSchemaBuilder; use datafusion_datasource::file_scan_config::FileScanConfigBuilder; -use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_expr::{ Operator, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility, @@ -59,6 +58,7 @@ use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeE use datafusion_physical_plan::streaming::{PartitionStream, StreamingTableExec}; use datafusion_physical_plan::union::UnionExec; use datafusion_physical_plan::{ExecutionPlan, displayable}; +use datafusion_storage::StorageUrl; use insta::assert_snapshot; use itertools::Itertools; @@ -385,20 +385,19 @@ fn create_simple_csv_exec() -> Arc { Field::new("d", DataType::Int32, true), Field::new("e", DataType::Int32, true), ])); - let config = - FileScanConfigBuilder::new(ObjectStoreUrl::parse("test:///").unwrap(), { - let options = CsvOptions { - has_header: Some(false), - delimiter: 0, - quote: 0, - ..Default::default() - }; - Arc::new(CsvSource::new(schema.clone()).with_csv_options(options)) - }) - .with_file(PartitionedFile::new("x", 100)) - .with_projection_indices(Some(vec![0, 1, 2, 3, 4])) - .unwrap() - .build(); + let config = FileScanConfigBuilder::new(StorageUrl::parse("test:///").unwrap(), { + let options = CsvOptions { + has_header: Some(false), + delimiter: 0, + quote: 0, + ..Default::default() + }; + Arc::new(CsvSource::new(schema.clone()).with_csv_options(options)) + }) + .with_file(PartitionedFile::new("x", 100)) + .with_projection_indices(Some(vec![0, 1, 2, 3, 4])) + .unwrap() + .build(); DataSourceExec::from_data_source(config) } @@ -410,20 +409,19 @@ fn create_projecting_csv_exec() -> Arc { Field::new("c", DataType::Int32, true), Field::new("d", DataType::Int32, true), ])); - let config = - FileScanConfigBuilder::new(ObjectStoreUrl::parse("test:///").unwrap(), { - let options = CsvOptions { - has_header: Some(false), - delimiter: 0, - quote: 0, - ..Default::default() - }; - Arc::new(CsvSource::new(schema.clone()).with_csv_options(options)) - }) - .with_file(PartitionedFile::new("x", 100)) - .with_projection_indices(Some(vec![3, 2, 1])) - .unwrap() - .build(); + let config = FileScanConfigBuilder::new(StorageUrl::parse("test:///").unwrap(), { + let options = CsvOptions { + has_header: Some(false), + delimiter: 0, + quote: 0, + ..Default::default() + }; + Arc::new(CsvSource::new(schema.clone()).with_csv_options(options)) + }) + .with_file(PartitionedFile::new("x", 100)) + .with_projection_indices(Some(vec![3, 2, 1])) + .unwrap() + .build(); DataSourceExec::from_data_source(config) } @@ -1584,7 +1582,7 @@ fn partitioned_data_source() -> Arc { ))]) .build(); let config = FileScanConfigBuilder::new( - ObjectStoreUrl::parse("test:///").unwrap(), + StorageUrl::parse("test:///").unwrap(), Arc::new(CsvSource::new(table_schema).with_csv_options(options)), ) .with_file(PartitionedFile::new("x", 100)) diff --git a/datafusion/core/tests/physical_optimizer/pushdown_utils.rs b/datafusion/core/tests/physical_optimizer/pushdown_utils.rs index 861ead100cc10..297a6b7669b67 100644 --- a/datafusion/core/tests/physical_optimizer/pushdown_utils.rs +++ b/datafusion/core/tests/physical_optimizer/pushdown_utils.rs @@ -17,7 +17,7 @@ use arrow::datatypes::SchemaRef; use arrow::{array::RecordBatch, compute::concat_batches}; -use datafusion::{datasource::object_store::ObjectStoreUrl, physical_plan::PhysicalExpr}; +use datafusion::{physical_plan::PhysicalExpr, storage::StorageUrl}; use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{Result, config::ConfigOptions, internal_err}; use datafusion_datasource::{ @@ -40,9 +40,9 @@ use datafusion_physical_plan::{ }, metrics::ExecutionPlanMetricsSet, }; +use datafusion_storage::StorageBinding; use futures::StreamExt; use futures::{FutureExt, Stream}; -use object_store::ObjectStore; use std::fmt::Write as _; use std::{ fmt::{Display, Formatter}, @@ -130,9 +130,10 @@ impl TestSource { impl FileSource for TestSource { fn create_file_opener( &self, - _object_store: Arc, + _object_store: Arc, _base_config: &FileScanConfig, _partition: usize, + _access_context: datafusion_storage::FileAccessContext, ) -> Result> { Ok(Arc::new(TestOpener { batches: self.batches.clone(), @@ -287,7 +288,7 @@ impl TestScanBuilder { self.batches, )); let base_config = - FileScanConfigBuilder::new(ObjectStoreUrl::parse("test://").unwrap(), source) + FileScanConfigBuilder::new(StorageUrl::parse("test://").unwrap(), source) .with_file(PartitionedFile::new("test.parquet", 123)) .build(); DataSourceExec::from_data_source(base_config) diff --git a/datafusion/core/tests/physical_optimizer/replace_with_order_preserving_variants.rs b/datafusion/core/tests/physical_optimizer/replace_with_order_preserving_variants.rs index b8adb283c2d31..c6c066b260944 100644 --- a/datafusion/core/tests/physical_optimizer/replace_with_order_preserving_variants.rs +++ b/datafusion/core/tests/physical_optimizer/replace_with_order_preserving_variants.rs @@ -128,10 +128,13 @@ impl ReplaceTest { ) .await .expect("could not create object store"); - ctx.register_object_store( + ctx.register_storage( &Url::parse("test://").unwrap(), - Arc::new(object_store), - ); + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + Arc::new(object_store), + )), + ) + .unwrap(); let task_ctx = Arc::new(TaskContext::from(&ctx)); let res = collect(optimized_physical_plan, task_ctx).await; assert!( diff --git a/datafusion/core/tests/physical_optimizer/test_utils.rs b/datafusion/core/tests/physical_optimizer/test_utils.rs index 6b10efffa90a3..7672f1242833b 100644 --- a/datafusion/core/tests/physical_optimizer/test_utils.rs +++ b/datafusion/core/tests/physical_optimizer/test_utils.rs @@ -38,7 +38,6 @@ use datafusion_common::{ ColumnStatistics, JoinType, NullEquality, Result, Statistics, internal_err, }; use datafusion_datasource::file_scan_config::FileScanConfigBuilder; -use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_expr::{WindowFrame, WindowFunctionDefinition}; use datafusion_functions_aggregate::count::count_udaf; @@ -74,11 +73,12 @@ use datafusion_physical_plan::{ InputDistributionRequirements, InputOrderMode, Partitioning, PlanProperties, ReplaceChildrenOptions, SortOrderPushdownResult, StatisticsArgs, displayable, }; +use datafusion_storage::StorageUrl; /// Create a non sorted parquet exec pub fn parquet_exec(schema: SchemaRef) -> Arc { let config = FileScanConfigBuilder::new( - ObjectStoreUrl::parse("test:///").unwrap(), + StorageUrl::parse("test:///").unwrap(), Arc::new(ParquetSource::new(schema)), ) .with_file(PartitionedFile::new("x".to_string(), 100)) @@ -93,7 +93,7 @@ pub(crate) fn parquet_exec_with_sort( output_ordering: Vec, ) -> Arc { let config = FileScanConfigBuilder::new( - ObjectStoreUrl::parse("test:///").unwrap(), + StorageUrl::parse("test:///").unwrap(), Arc::new(ParquetSource::new(schema)), ) .with_file(PartitionedFile::new("x".to_string(), 100)) @@ -131,7 +131,7 @@ pub(crate) fn parquet_exec_with_stats(file_size: u64) -> Arc { statistics.column_statistics = column_stats(); let config = FileScanConfigBuilder::new( - ObjectStoreUrl::parse("test:///").unwrap(), + StorageUrl::parse("test:///").unwrap(), Arc::new(ParquetSource::new(schema())), ) .with_file(PartitionedFile::new("x".to_string(), file_size)) diff --git a/datafusion/core/tests/sql/explain_analyze.rs b/datafusion/core/tests/sql/explain_analyze.rs index 2ca04ad6b3fac..616466496b12b 100644 --- a/datafusion/core/tests/sql/explain_analyze.rs +++ b/datafusion/core/tests/sql/explain_analyze.rs @@ -23,7 +23,7 @@ use datafusion::config::ConfigOptions; use datafusion::physical_plan::display::DisplayableExecutionPlan; use datafusion::physical_plan::metrics::Timestamp; use datafusion_common::format::{ExplainAnalyzeCategories, MetricCategory, MetricType}; -use object_store::path::Path; +use datafusion_storage::path::Path; #[tokio::test] async fn explain_analyze_baseline_metrics() { diff --git a/datafusion/core/tests/sql/mod.rs b/datafusion/core/tests/sql/mod.rs index afed2f82d57a8..51c045a97d852 100644 --- a/datafusion/core/tests/sql/mod.rs +++ b/datafusion/core/tests/sql/mod.rs @@ -34,7 +34,7 @@ use datafusion::{execution::context::SessionContext, physical_plan::displayable} use datafusion_common::test_util::batches_to_sort_string; use datafusion_common::utils::get_available_parallelism; use datafusion_common::{assert_contains, assert_not_contains}; -use object_store::path::Path; +use datafusion_storage::path::Path; use std::fs::File; use std::io::Write; use std::path::PathBuf; diff --git a/datafusion/core/tests/sql/path_partition.rs b/datafusion/core/tests/sql/path_partition.rs index f8701ccc568a3..7ce9f279a682f 100644 --- a/datafusion/core/tests/sql/path_partition.rs +++ b/datafusion/core/tests/sql/path_partition.rs @@ -47,11 +47,12 @@ use futures::StreamExt; use futures::stream::{self, BoxStream}; use insta::assert_snapshot; use object_store::{ - Attributes, CopyOptions, GetRange, MultipartUpload, PutMultipartOptions, PutPayload, + Attributes, CopyOptions, GetRange, MultipartUpload, ObjectMeta, ObjectStore, + PutMultipartOptions, PutPayload, }; use object_store::{ - GetOptions, GetResult, GetResultPayload, ListResult, ObjectMeta, ObjectStore, - PutOptions, PutResult, path::Path, + GetOptions, GetResult, GetResultPayload, ListResult, PutOptions, PutResult, + path::Path, }; use url::Url; @@ -547,10 +548,13 @@ fn register_partitioned_aggregate_csv( let csv_file_path = format!("{testdata}/csv/aggregate_test_100.csv"); let file_schema = test_util::aggr_test_schema(); let url = Url::parse("mirror://").unwrap(); - ctx.register_object_store( + ctx.register_storage( &url, - MirroringObjectStore::new_arc(csv_file_path, store_paths), - ); + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + MirroringObjectStore::new_arc(csv_file_path, store_paths), + )), + ) + .unwrap(); let options = ListingOptions::new(Arc::new(CsvFormat::default())) .with_table_partition_cols( @@ -599,10 +603,13 @@ async fn create_partitioned_alltypes_parquet_table( let testdata = parquet_test_data(); let parquet_file_path = format!("{testdata}/{source_file}"); let url = Url::parse("mirror://").unwrap(); - ctx.register_object_store( + ctx.register_storage( &url, - MirroringObjectStore::new_arc(parquet_file_path.clone(), store_paths), - ); + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + MirroringObjectStore::new_arc(parquet_file_path.clone(), store_paths), + )), + ) + .unwrap(); let options = ListingOptions::new(Arc::new(ParquetFormat::default())) .with_table_partition_cols( diff --git a/datafusion/core/tests/storage.rs b/datafusion/core/tests/storage.rs new file mode 100644 index 0000000000000..66739d71f8efb --- /dev/null +++ b/datafusion/core/tests/storage.rs @@ -0,0 +1,512 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#![cfg(all(feature = "parquet", feature = "sql", feature = "opendal"))] +use async_trait::async_trait; +use bytes::Bytes; +use datafusion::datasource::file_format::options::{ArrowReadOptions, JsonReadOptions}; +use datafusion::{ + arrow::{ + array::Int32Array, + datatypes::{DataType, Field, Schema}, + record_batch::RecordBatch, + }, + prelude::*, +}; +use datafusion_common::{Result, assert_batches_eq}; +use datafusion_datasource::PartitionedFile; +use datafusion_datasource_parquet::{ + DefaultParquetFileReaderFactory, ParquetFileReaderFactory, + storage::StorageParquetTable, +}; +use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; +use datafusion_storage::{ + FileAccessContext, FileInfo, FileReader, ReadRange, Storage, StorageUrl, + WriterOptions, +}; +use futures::{StreamExt, stream::BoxStream}; +use parquet::{ + arrow::{ArrowWriter, async_reader::AsyncFileReader}, + file::properties::WriterProperties, +}; +use std::{ + ops::Range, + sync::{ + Arc, Mutex, + atomic::{AtomicUsize, Ordering}, + }, +}; +use tokio::io::AsyncWriteExt; + +fn parquet_bytes(start: i32) -> Bytes { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("value", DataType::Int32, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from_iter_values(start..start + 12)), + Arc::new(Int32Array::from_iter_values( + (start..start + 12).map(|i| i * 10), + )), + ], + ) + .unwrap(); + let mut bytes = Vec::new(); + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(3)) + .build(); + let mut writer = ArrowWriter::try_new(&mut bytes, schema, Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + Bytes::from(bytes) +} + +#[tokio::test] +async fn opendal_discovery_sql_and_atomic_replacement() -> Result<()> { + let ctx = SessionContext::new(); + let url = url::Url::parse("memory://warehouse/events/").unwrap(); + let operator = opendal::Operator::new(opendal::services::Memory::default()).unwrap(); + ctx.register_storage( + &url, + Arc::new(datafusion_storage_opendal::OpendalStorage::new(operator)), + ) + .unwrap(); + let bound = ctx.runtime_env().storage_registry.get(&url).unwrap(); + // The same registration supports output creation and subsequent discovery. + let data = parquet_bytes(0); + let file = FileInfo::new( + datafusion_storage::path::Path::from("events/data.parquet"), + data.len() as u64, + ); + let mut writer = bound + .writer( + &file.location, + WriterOptions::default(), + FileAccessContext::new("write"), + ) + .await + .unwrap(); + writer.write_all(&data).await.unwrap(); + writer.shutdown().await.unwrap(); + let old = ctx + .read_parquet(url.as_str(), ParquetReadOptions::default()) + .await?; + ctx.register_table("old", old.clone().into_view())?; + let plan = ctx + .sql("SELECT value FROM old WHERE id >= 9 ORDER BY value") + .await? + .create_physical_plan() + .await?; + let plan_text = datafusion_physical_plan::displayable(plan.as_ref()) + .indent(true) + .to_string(); + assert!( + plan_text.contains("file_type=parquet") + && plan_text.contains("pruning_predicate="), + "{plan_text}" + ); + let result = ctx + .sql("SELECT value FROM old WHERE id >= 9 ORDER BY value") + .await? + .collect() + .await?; + assert_batches_eq!( + [ + "+-------+", + "| value |", + "+-------+", + "| 90 |", + "| 100 |", + "| 110 |", + "+-------+" + ], + &result + ); + + let replacement = + opendal::Operator::new(opendal::services::Memory::default()).unwrap(); + replacement + .write("events/data.parquet", parquet_bytes(100)) + .await + .unwrap(); + ctx.register_storage( + &url, + Arc::new(datafusion_storage_opendal::OpendalStorage::new(replacement)), + ) + .unwrap(); + let new = ctx + .read_parquet(url.as_str(), ParquetReadOptions::default()) + .await?; + ctx.deregister_storage(&url).unwrap(); + assert!( + ctx.read_parquet(url.as_str(), ParquetReadOptions::default()) + .await + .is_err() + ); + // Both tables retain their original complete bindings after deregistration. + assert_eq!( + old.collect().await?[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 0 + ); + assert_eq!( + new.collect().await?[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 100 + ); + Ok(()) +} + +#[derive(Debug)] +struct ReadOnly { + bytes: Bytes, + contexts: Arc>>, + ranges: Arc, +} +#[async_trait] +impl Storage for ReadOnly { + async fn open( + &self, + _: &datafusion_storage::path::Path, + context: FileAccessContext, + ) -> datafusion_storage::Result> { + self.contexts.lock().unwrap().push(context); + Ok(Arc::new(ReadOnly { + bytes: self.bytes.clone(), + contexts: Arc::clone(&self.contexts), + ranges: Arc::clone(&self.ranges), + })) + } +} +#[async_trait] +impl FileReader for ReadOnly { + async fn read_range(&self, range: ReadRange) -> datafusion_storage::Result { + let range = match range { + ReadRange::Bounded(r) => r, + ReadRange::Suffix(n) => { + let size = self.bytes.len() as u64; + size.saturating_sub(n)..size + } + }; + self.ranges.fetch_add(1, Ordering::Relaxed); + Ok(self.bytes.slice(range.start as usize..range.end as usize)) + } + + async fn read_ranges( + &self, + ranges: Vec>, + ) -> datafusion_storage::Result> { + self.ranges.fetch_add(ranges.len(), Ordering::Relaxed); + Ok(ranges + .into_iter() + .map(|r| self.bytes.slice(r.start as usize..r.end as usize)) + .collect()) + } + fn stream( + self: Arc, + _: Option, + ) -> BoxStream<'static, datafusion_storage::Result> { + // Parquet must use the range API, without requiring a stream implementation. + futures::stream::once(async { + Err(datafusion_storage::Error::NotSupported("stream".into())) + }) + .boxed() + } +} +#[tokio::test] +async fn read_only_manifest_reuses_native_scan_with_fresh_execution_contexts() +-> Result<()> { + let ctx = SessionContext::new(); + let url = StorageUrl::parse("custom://manifest").unwrap(); + let bytes = parquet_bytes(0); + let file = FileInfo::new( + datafusion_storage::path::Path::from("part.parquet"), + bytes.len() as u64, + ); + let contexts = Arc::new(Mutex::new(Vec::new())); + let ranges = Arc::new(AtomicUsize::new(0)); + let storage: Arc = Arc::new(ReadOnly { + bytes, + contexts: Arc::clone(&contexts), + ranges: Arc::clone(&ranges), + }); + ctx.register_storage(url.as_ref(), storage).unwrap(); + let binding = ctx + .runtime_env() + .storage_registry + .get(url.as_ref()) + .unwrap(); + let table = StorageParquetTable::try_new(binding, vec![file], None).await?; + ctx.register_table("manifest", Arc::new(table))?; + let query = ctx + .sql("SELECT value FROM manifest WHERE id > 8 ORDER BY value") + .await?; + let first = query.clone().collect().await?; + // Cancelling a completed execution must not poison a later execution of its table. + for context in contexts.lock().unwrap().iter() { + context.cancellation.cancel(); + } + ctx.deregister_storage(url.as_ref()).unwrap(); + let second = query.collect().await?; + assert_eq!(first, second); + assert!(contexts.lock().unwrap().len() >= 3); // inference and two executions + assert!(ranges.load(Ordering::Relaxed) > 3); + assert_batches_eq!( + [ + "+-------+", + "| value |", + "+-------+", + "| 90 |", + "| 100 |", + "| 110 |", + "+-------+" + ], + &second + ); + Ok(()) +} + +#[derive(Debug)] +struct CustomFactory { + inner: DefaultParquetFileReaderFactory, + calls: Arc, +} +impl ParquetFileReaderFactory for CustomFactory { + fn create_reader( + &self, + partition: usize, + file: PartitionedFile, + hint: Option, + metrics: &ExecutionPlanMetricsSet, + ) -> Result> { + self.calls.fetch_add(1, Ordering::Relaxed); + self.inner.create_reader(partition, file, hint, metrics) + } +} +#[tokio::test] +async fn custom_parquet_factory_is_preserved_without_object_store_registration() +-> Result<()> { + let ctx = SessionContext::new(); + let url = StorageUrl::parse("custom://factory").unwrap(); + let bytes = parquet_bytes(0); + let file = FileInfo::new( + datafusion_storage::path::Path::from("part.parquet"), + bytes.len() as u64, + ); + let storage: Arc = Arc::new(ReadOnly { + bytes, + contexts: Arc::new(Mutex::new(Vec::new())), + ranges: Arc::new(AtomicUsize::new(0)), + }); + let calls = Arc::new(AtomicUsize::new(0)); + let factory = Arc::new(CustomFactory { + inner: DefaultParquetFileReaderFactory::new(Arc::new( + datafusion_storage::StorageBinding::new(url.clone(), Arc::clone(&storage)), + )), + calls: Arc::clone(&calls), + }); + let table = StorageParquetTable::try_new( + Arc::new(datafusion_storage::StorageBinding::new( + url.clone(), + storage, + )), + vec![file.clone()], + None, + ) + .await? + .with_parquet_file_reader_factory(factory.clone()); + assert_eq!( + ctx.read_table(Arc::new(table))? + .collect() + .await? + .iter() + .map(RecordBatch::num_rows) + .sum::(), + 12 + ); + { + // The legacy format planning hook must preserve an explicitly supplied factory too. + use datafusion_datasource::{ + file::FileSource, file_format::FileFormat, + file_scan_config::FileScanConfigBuilder, + }; + use datafusion_datasource_parquet::{ParquetFormat, source::ParquetSource}; + let source: Arc = Arc::new( + ParquetSource::new(Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("value", DataType::Int32, false), + ]))) + .with_parquet_file_reader_factory(factory), + ); + let config = FileScanConfigBuilder::new(StorageUrl::parse(url.as_str())?, source) + .with_file(PartitionedFile::new(file.location.to_string(), file.size)) + .build(); + let plan = ParquetFormat::default() + .create_physical_plan(&ctx.state(), config) + .await?; + assert_eq!( + datafusion_physical_plan::collect(plan, ctx.task_ctx()) + .await? + .iter() + .map(RecordBatch::num_rows) + .sum::(), + 12 + ); + } + assert!(calls.load(Ordering::Relaxed) > 0); + Ok(()) +} + +#[tokio::test] +async fn dropping_a_scan_cancels_its_storage_context() -> Result<()> { + let ctx = SessionContext::new_with_config(SessionConfig::new().with_batch_size(2)); + let url = StorageUrl::parse("custom://cancel").unwrap(); + let bytes = parquet_bytes(0); + let file = FileInfo::new( + datafusion_storage::path::Path::from("part.parquet"), + bytes.len() as u64, + ); + let contexts = Arc::new(Mutex::new(Vec::new())); + let storage: Arc = Arc::new(ReadOnly { + bytes, + contexts: Arc::clone(&contexts), + ranges: Arc::new(AtomicUsize::new(0)), + }); + let table = StorageParquetTable::try_new( + Arc::new(datafusion_storage::StorageBinding::new(url, storage)), + vec![file], + None, + ) + .await?; + let mut stream = ctx.read_table(Arc::new(table))?.execute_stream().await?; + assert_eq!(stream.next().await.unwrap()?.num_rows(), 2); + drop(stream); + assert!( + contexts + .lock() + .unwrap() + .iter() + .all(|context| context.cancellation.is_cancelled()) + ); + Ok(()) +} + +#[tokio::test] +async fn one_opendal_registration_supports_builtin_formats_and_copy() -> Result<()> { + let ctx = SessionContext::new(); + let url = url::Url::parse("memory://formats/").unwrap(); + let operator = opendal::Operator::new(opendal::services::Memory::default()).unwrap(); + ctx.register_storage( + &url, + Arc::new(datafusion_storage_opendal::OpendalStorage::new(operator)), + )?; + for format in ["CSV", "JSON", "ARROW", "PARQUET"] { + let path = format!("memory://formats/{}/", format.to_lowercase()); + let options = if format == "CSV" { + "OPTIONS ('format.has_header' 'true')" + } else { + "" + }; + ctx.sql(&format!( + "COPY (SELECT * FROM (VALUES (1), (2), (3)) AS t(id)) TO '{path}' STORED AS {format} {options}" + )).await?.collect().await?; + let frame = match format { + "CSV" => { + ctx.read_csv(&path, CsvReadOptions::new().has_header(true)) + .await? + } + "JSON" => ctx.read_json(&path, JsonReadOptions::default()).await?, + "ARROW" => ctx.read_arrow(&path, ArrowReadOptions::default()).await?, + "PARQUET" => { + ctx.read_parquet(&path, ParquetReadOptions::default()) + .await? + } + _ => unreachable!(), + }; + ctx.register_table(format.to_lowercase(), frame.into_view())?; + let batches = ctx + .sql(&format!( + "SELECT SUM(id) AS total FROM {}", + format.to_lowercase() + )) + .await? + .collect() + .await?; + assert_batches_eq!( + [ + "+-------+", + "| total |", + "+-------+", + "| 6 |", + "+-------+" + ], + &batches + ); + } + Ok(()) +} + +#[cfg(feature = "avro")] +#[tokio::test] +async fn opendal_reads_avro_through_the_normal_api() -> Result<()> { + let operator = opendal::Operator::new(opendal::services::Memory::default()).unwrap(); + let data = std::fs::read(concat!( + env!("CARGO_MANIFEST_DIR"), + "/../../testing/data/avro/alltypes_plain.avro" + ))?; + operator.write("data.avro", data).await.unwrap(); + let ctx = SessionContext::new(); + ctx.register_storage( + &url::Url::parse("memory://avro").unwrap(), + Arc::new(datafusion_storage_opendal::OpendalStorage::new(operator)), + )?; + assert_eq!( + ctx.read_avro("memory://avro/data.avro", AvroReadOptions::default()) + .await? + .count() + .await?, + 8 + ); + Ok(()) +} + +#[tokio::test] +async fn opendal_listing_table_insert_uses_the_registered_writer() -> Result<()> { + let operator = opendal::Operator::new(opendal::services::Memory::default()).unwrap(); + let ctx = SessionContext::new(); + ctx.register_storage( + &url::Url::parse("memory://insert").unwrap(), + Arc::new(datafusion_storage_opendal::OpendalStorage::new(operator)), + )?; + ctx.sql("CREATE EXTERNAL TABLE events (id INT) STORED AS PARQUET LOCATION 'memory://insert/events/'").await?.collect().await?; + ctx.sql("INSERT INTO events VALUES (1), (2), (3)") + .await? + .collect() + .await?; + assert_eq!(ctx.table("events").await?.count().await?, 3); + Ok(()) +} diff --git a/datafusion/core/tests/tracing/mod.rs b/datafusion/core/tests/tracing/mod.rs index 9f72c2c66f4cd..a50cabb0da7e2 100644 --- a/datafusion/core/tests/tracing/mod.rs +++ b/datafusion/core/tests/tracing/mod.rs @@ -97,7 +97,13 @@ async fn run_query() { // Register the traceable object store with a test URL. let url = Url::parse("test://").unwrap(); - ctx.register_object_store(&url, traceable_store.clone()); + ctx.register_storage( + &url, + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + traceable_store.clone(), + )), + ) + .unwrap(); // Register a listing table from the test data directory. let table_path = format!("test://{test_data}/"); diff --git a/datafusion/datasource-arrow/Cargo.toml b/datafusion/datasource-arrow/Cargo.toml index 6f50135403d69..31ff02b4f69a0 100644 --- a/datafusion/datasource-arrow/Cargo.toml +++ b/datafusion/datasource-arrow/Cargo.toml @@ -35,7 +35,7 @@ arrow = { workspace = true } arrow-ipc = { workspace = true } async-trait = { workspace = true } bytes = { workspace = true } -datafusion-common = { workspace = true, features = ["object_store"] } +datafusion-common = { workspace = true } datafusion-common-runtime = { workspace = true } datafusion-datasource = { workspace = true } datafusion-execution = { workspace = true } @@ -44,13 +44,15 @@ datafusion-physical-expr-common = { workspace = true } datafusion-physical-plan = { workspace = true } datafusion-proto-models = { workspace = true, optional = true } datafusion-session = { workspace = true } +datafusion-storage.workspace = true futures = { workspace = true } itertools = { workspace = true } -object_store = { workspace = true } tokio = { workspace = true } [dev-dependencies] chrono = { workspace = true } +object_store = { workspace = true } +test-utils = { path = "../../test-utils" } # Note: add additional linter rules in lib.rs. # Rust does not support workspace + new linter rules in subcrates yet diff --git a/datafusion/datasource-arrow/src/file_format.rs b/datafusion/datasource-arrow/src/file_format.rs index 2bee57ef17581..6dd24039564f2 100644 --- a/datafusion/datasource-arrow/src/file_format.rs +++ b/datafusion/datasource-arrow/src/file_format.rs @@ -21,13 +21,11 @@ use std::collections::HashMap; use std::fmt::{self, Debug}; -use std::io::{Seek, SeekFrom}; use std::sync::Arc; use arrow::datatypes::{Schema, SchemaRef}; use arrow::error::ArrowError; use arrow::ipc::convert::fb_to_schema; -use arrow::ipc::reader::{FileReader, StreamReader}; use arrow::ipc::writer::IpcWriteOptions; use arrow::ipc::{CompressionType, root_as_message}; use datafusion_common::error::Result; @@ -41,9 +39,7 @@ use datafusion_datasource::display::FileGroupDisplay; use datafusion_datasource::file::FileSource; use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; use datafusion_datasource::sink::{DataSink, DataSinkExec}; -use datafusion_datasource::write::{ - ObjectWriterBuilder, SharedBuffer, get_writer_schema, -}; +use datafusion_datasource::write::{FileWriterBuilder, SharedBuffer, get_writer_schema}; use datafusion_datasource::{TableSchema, TableSchemaBuilder}; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_expr::dml::InsertOp; @@ -59,12 +55,10 @@ use datafusion_datasource::source::DataSourceExec; use datafusion_datasource::write::demux::DemuxedStreamReceiver; use datafusion_physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan}; use datafusion_session::Session; +use datafusion_storage::{FileAccessContext, path::Path}; +use datafusion_storage::{FileInfo, StorageBinding}; use futures::StreamExt; use futures::stream::BoxStream; -use object_store::{ - GetOptions, GetRange, GetResultPayload, ObjectMeta, ObjectStore, ObjectStoreExt, - path::Path, -}; use tokio::io::AsyncWriteExt; /// Initial writing buffer size. Note this is just a size hint for efficiency. It @@ -136,35 +130,13 @@ impl FileFormat for ArrowFormat { async fn infer_schema( &self, _state: &dyn Session, - store: &Arc, - objects: &[ObjectMeta], + store: &Arc, + objects: &[FileInfo], ) -> Result { let mut schemas = vec![]; for object in objects { - let r = store.as_ref().get(&object.location).await?; - let schema = match r.payload { - #[cfg(not(target_arch = "wasm32"))] - GetResultPayload::File(mut file, _) => { - match FileReader::try_new(&mut file, None) { - Ok(reader) => reader.schema(), - Err(file_error) => { - // not in the file format, but FileReader read some bytes - // while trying to parse the file and so we need to rewind - // it to the beginning of the file - file.seek(SeekFrom::Start(0))?; - match StreamReader::try_new(&mut file, None) { - Ok(reader) => reader.schema(), - Err(stream_error) => { - return Err(internal_datafusion_err!( - "Failed to parse Arrow file as either file format or stream format. File format error: {file_error}. Stream format error: {stream_error}" - )); - } - } - } - } - } - GetResultPayload::Stream(stream) => infer_stream_schema(stream).await?, - }; + let reader = store.open(object, FileAccessContext::new("schema")).await?; + let schema = infer_stream_schema(reader.stream(None)).await?; schemas.push(Arc::unwrap_or_clone(schema)); } let merged_schema = Schema::try_merge(schemas)?; @@ -174,9 +146,9 @@ impl FileFormat for ArrowFormat { async fn infer_stats( &self, _state: &dyn Session, - _store: &Arc, + _store: &Arc, table_schema: SchemaRef, - _object: &ObjectMeta, + _object: &FileInfo, ) -> Result { Ok(Statistics::new_unknown(&table_schema)) } @@ -186,7 +158,7 @@ impl FileFormat for ArrowFormat { state: &dyn Session, conf: FileScanConfig, ) -> Result> { - let object_store = state.runtime_env().object_store(&conf.object_store_url)?; + let object_store = state.runtime_env().storage(&conf.object_store_url)?; let object_location = &conf .file_groups .first() @@ -266,7 +238,7 @@ impl FileSink for ArrowFileSink { context: &Arc, demux_task: SpawnedTask>, mut file_stream_rx: DemuxedStreamReceiver, - object_store: Arc, + object_store: Arc, ) -> Result { let mut file_write_tasks: JoinSet> = JoinSet::new(); @@ -281,7 +253,7 @@ impl FileSink for ArrowFileSink { &get_writer_schema(&self.config), ipc_options.clone(), )?; - let mut object_store_writer = ObjectWriterBuilder::new( + let mut object_store_writer = FileWriterBuilder::new( FileCompressionType::UNCOMPRESSED, &path, Arc::clone(&object_store), @@ -293,7 +265,8 @@ impl FileSink for ArrowFileSink { .execution .objectstore_writer_buffer_size, )) - .build()?; + .build() + .await?; file_write_tasks.spawn(async move { let mut row_count = 0; while let Some(batch) = rx.recv().await { @@ -384,7 +357,7 @@ const ARROW_MAGIC: [u8; 6] = *b"ARROW1"; const CONTINUATION_MARKER: [u8; 4] = [0xff; 4]; async fn infer_stream_schema( - mut stream: BoxStream<'static, object_store::Result>, + mut stream: BoxStream<'static, datafusion_storage::Result>, ) -> Result { // IPC streaming format. // See https://arrow.apache.org/docs/format/Columnar.html#ipc-streaming-format @@ -426,17 +399,25 @@ async fn infer_stream_schema( // For the purposes of this function, the arrow "preamble" is the magic number, padding, // and the continuation marker. 16 bytes covers the preamble and metadata length // no matter which version or format is used. - let bytes = extend_bytes_to_n_length_from_stream(vec![], 16, &mut stream).await?; + let mut bytes = extend_bytes_to_n_length_from_stream(vec![], 16, &mut stream).await?; // The preamble length is everything before the metadata length let preamble_len = if bytes[0..6] == ARROW_MAGIC { - // File format starts with magic number "ARROW1" - if bytes[8..12] == CONTINUATION_MARKER { - // Continuation marker was added in v0.15.0 - 12 + // IPC writers may align the first message beyond the minimum 8-byte + // prefix. Skip zero padding before reading its length/continuation marker. + let mut offset = 8; + loop { + bytes = extend_bytes_to_n_length_from_stream(bytes, offset + 4, &mut stream) + .await?; + if bytes[offset..offset + 4] != [0; 4] { + break; + } + offset += 4; + } + if bytes[offset..offset + 4] == CONTINUATION_MARKER { + offset + 4 } else { - // File format before v0.15.0 - 8 + offset } } else if bytes[0..4] == CONTINUATION_MARKER { // Stream format after v0.15.0 starts with continuation marker @@ -446,6 +427,8 @@ async fn infer_stream_schema( 0 }; + bytes = extend_bytes_to_n_length_from_stream(bytes, preamble_len + 4, &mut stream) + .await?; let meta_len_bytes: [u8; 4] = bytes[preamble_len..preamble_len + 4] .try_into() .map_err(|err| { @@ -489,7 +472,7 @@ async fn infer_stream_schema( async fn extend_bytes_to_n_length_from_stream( bytes: Vec, n: usize, - stream: &mut BoxStream<'static, object_store::Result>, + stream: &mut BoxStream<'static, datafusion_storage::Result>, ) -> Result> { if bytes.len() >= n { return Ok(bytes); @@ -516,24 +499,23 @@ async fn extend_bytes_to_n_length_from_stream( } async fn is_object_in_arrow_ipc_file_format( - store: Arc, + store: Arc, object_location: &Path, ) -> Result { - let get_opts = GetOptions { - range: Some(GetRange::Bounded(0..6)), - ..Default::default() - }; - let bytes = store - .get_opts(object_location, get_opts) - .await? - .bytes() - .await?; + let context = FileAccessContext::new("format"); + let info = store.stat(object_location, &context).await?; + if info.size == 0 { + return Err(ArrowError::ParseError("Empty Arrow IPC file".into()).into()); + } + let reader = store.open(&info, context).await?; + let bytes = reader.read_range((0..info.size.min(6)).into()).await?; Ok(bytes.len() >= 6 && bytes[0..6] == ARROW_MAGIC) } #[cfg(test)] mod tests { use super::*; + use object_store::{ObjectStore, ObjectStoreExt}; use std::any::Any; @@ -646,10 +628,16 @@ mod tests { bytes.truncate(bytes.len() - 20); // mangle end to show we don't need to read whole file let location = Path::parse(file)?; let in_memory_store: Arc = Arc::new(InMemory::new()); - in_memory_store.put(&location, bytes.into()).await?; + in_memory_store + .put( + &object_store::path::Path::from(location.as_ref()), + bytes.into(), + ) + .await + .unwrap(); let state = MockSession::new(); - let object_meta = ObjectMeta { + let object_meta = FileInfo { location, last_modified: DateTime::default(), size: u64::MAX, @@ -668,7 +656,7 @@ mod tests { let inferred_schema = arrow_format .infer_schema( &state, - &(store.clone() as Arc), + &test_utils::storage::object_store(store.clone()), std::slice::from_ref(&object_meta), ) .await?; @@ -690,13 +678,19 @@ mod tests { bytes.truncate(20); // should cause error that file shorter than expected let location = Path::parse(file)?; let in_memory_store: Arc = Arc::new(InMemory::new()); - in_memory_store.put(&location, bytes.into()).await?; + in_memory_store + .put( + &object_store::path::Path::from(location.as_ref()), + bytes.into(), + ) + .await + .unwrap(); let state = MockSession::new(); - let object_meta = ObjectMeta { + let object_meta = FileInfo { location, last_modified: DateTime::default(), - size: u64::MAX, + size: 20, e_tag: None, version: None, }; @@ -707,7 +701,7 @@ mod tests { let err = arrow_format .infer_schema( &state, - &(store.clone() as Arc), + &test_utils::storage::object_store(store.clone()), std::slice::from_ref(&object_meta), ) .await; @@ -728,9 +722,19 @@ mod tests { let path = Path::from("test.arrow"); let file_bytes = std::fs::read("tests/data/example.arrow")?; - store.put(&path, file_bytes.into()).await?; + store + .put( + &object_store::path::Path::from(path.as_ref()), + file_bytes.into(), + ) + .await + .unwrap(); - let is_file = is_object_in_arrow_ipc_file_format(store.clone(), &path).await?; + let is_file = is_object_in_arrow_ipc_file_format( + test_utils::storage::object_store(store.clone()), + &path, + ) + .await?; assert!(is_file, "Should detect file format"); Ok(()) } @@ -741,9 +745,19 @@ mod tests { let path = Path::from("test_stream.arrow"); let stream_bytes = std::fs::read("tests/data/example_stream.arrow")?; - store.put(&path, stream_bytes.into()).await?; + store + .put( + &object_store::path::Path::from(path.as_ref()), + stream_bytes.into(), + ) + .await + .unwrap(); - let is_file = is_object_in_arrow_ipc_file_format(store.clone(), &path).await?; + let is_file = is_object_in_arrow_ipc_file_format( + test_utils::storage::object_store(store.clone()), + &path, + ) + .await?; assert!(!is_file, "Should detect stream format (not file)"); @@ -756,10 +770,18 @@ mod tests { let path = Path::from("corrupted.arrow"); store - .put(&path, Bytes::from(vec![0x43, 0x4f, 0x52, 0x41]).into()) - .await?; + .put( + &object_store::path::Path::from(path.as_ref()), + Bytes::from(vec![0x43, 0x4f, 0x52, 0x41]).into(), + ) + .await + .unwrap(); - let is_file = is_object_in_arrow_ipc_file_format(store.clone(), &path).await?; + let is_file = is_object_in_arrow_ipc_file_format( + test_utils::storage::object_store(store.clone()), + &path, + ) + .await?; assert!( !is_file, @@ -774,11 +796,21 @@ mod tests { let store = Arc::new(InMemory::new()); let path = Path::from("empty.arrow"); - store.put(&path, Bytes::new().into()).await?; + store + .put( + &object_store::path::Path::from(path.as_ref()), + Bytes::new().into(), + ) + .await + .unwrap(); - let result = is_object_in_arrow_ipc_file_format(store.clone(), &path).await; + let result = is_object_in_arrow_ipc_file_format( + test_utils::storage::object_store(store.clone()), + &path, + ) + .await; - // currently errors because it tries to read 0..6 from an empty file + // Empty data cannot identify either Arrow IPC format. assert!(result.is_err(), "Empty file should error"); Ok(()) diff --git a/datafusion/datasource-arrow/src/source.rs b/datafusion/datasource-arrow/src/source.rs index 7e8189a78d392..98f59662dcad5 100644 --- a/datafusion/datasource-arrow/src/source.rs +++ b/datafusion/datasource-arrow/src/source.rs @@ -51,9 +51,10 @@ use datafusion_physical_plan::projection::ProjectionExprs; use datafusion_datasource::file_stream::FileOpenFuture; use datafusion_datasource::file_stream::FileOpener; +use datafusion_storage::FileAccessContext; +use datafusion_storage::StorageBinding; use futures::StreamExt; use itertools::Itertools; -use object_store::{GetOptions, GetRange, GetResultPayload, ObjectStore, ObjectStoreExt}; /// Enum indicating which Arrow IPC format to use #[derive(Clone, Copy, Debug)] @@ -64,9 +65,29 @@ enum ArrowFormat { Stream, } +/// Read the IPC file footer independently of the writer's message alignment. +pub(crate) async fn read_file_footer( + reader: &dyn datafusion_storage::FileReader, +) -> Result { + use datafusion_storage::ReadRange; + let trailer = reader.read_range(ReadRange::Suffix(10)).await?; + let trailer: [u8; 10] = trailer.as_ref().try_into().map_err(|_| { + exec_datafusion_err!("Arrow IPC file is shorter than its trailer") + })?; + let footer_len = arrow_ipc::reader::read_footer_length(trailer)? as u64; + let bytes = reader + .read_range(ReadRange::Suffix(10 + footer_len)) + .await?; + if bytes.len() < 10 + footer_len as usize { + return Err(exec_datafusion_err!("Arrow IPC footer exceeds file size")); + } + Ok(bytes.slice(..bytes.len() - 10)) +} + /// `FileOpener` for Arrow IPC stream format. Supports only sequential reading. pub(crate) struct ArrowStreamFileOpener { - object_store: Arc, + object_store: Arc, + access_context: FileAccessContext, projection: Option>, } @@ -78,30 +99,21 @@ impl FileOpener for ArrowStreamFileOpener { )); } let object_store = Arc::clone(&self.object_store); + let access_context = self.access_context.clone(); let projection = self.projection.clone(); Ok(Box::pin(async move { - let r = object_store - .get(&partitioned_file.object_meta.location) + let reader = object_store + .open(&partitioned_file.object_meta, access_context) .await?; - let stream = match r.payload { - #[cfg(not(target_arch = "wasm32"))] - GetResultPayload::File(file, _) => futures::stream::iter( - StreamReader::try_new(file.try_clone()?, projection.clone())?, - ) - .map(|r| r.map_err(Into::into)) - .boxed(), - GetResultPayload::Stream(_) => { - let bytes = r.bytes().await?; - let cursor = Cursor::new(bytes); - futures::stream::iter(StreamReader::try_new( - cursor, - projection.clone(), - )?) + let stream = { + let bytes = + datafusion_storage::collect_bytes(reader.stream(None)).await?; + let cursor = Cursor::new(bytes); + futures::stream::iter(StreamReader::try_new(cursor, projection.clone())?) .map(|r| r.map_err(Into::into)) .boxed() - } }; Ok(stream) @@ -111,72 +123,46 @@ impl FileOpener for ArrowStreamFileOpener { /// `FileOpener` for Arrow IPC file format. Supports range-based parallel reading. pub(crate) struct ArrowFileOpener { - object_store: Arc, + object_store: Arc, + access_context: FileAccessContext, projection: Option>, } impl FileOpener for ArrowFileOpener { fn open(&self, partitioned_file: PartitionedFile) -> Result { let object_store = Arc::clone(&self.object_store); + let access_context = self.access_context.clone(); let projection = self.projection.clone(); Ok(Box::pin(async move { + let reader = object_store + .open(&partitioned_file.object_meta, access_context) + .await?; let range = partitioned_file.range.clone(); match range { None => { - let r = object_store - .get(&partitioned_file.object_meta.location) - .await?; - let stream = match r.payload { - #[cfg(not(target_arch = "wasm32"))] - GetResultPayload::File(file, _) => futures::stream::iter( - FileReader::try_new(file.try_clone()?, projection.clone())?, - ) + let stream = { + let bytes = + datafusion_storage::collect_bytes(reader.stream(None)) + .await?; + let cursor = Cursor::new(bytes); + futures::stream::iter(FileReader::try_new( + cursor, + projection.clone(), + )?) .map(|r| r.map_err(Into::into)) - .boxed(), - GetResultPayload::Stream(_) => { - let bytes = r.bytes().await?; - let cursor = Cursor::new(bytes); - futures::stream::iter(FileReader::try_new( - cursor, - projection.clone(), - )?) - .map(|r| r.map_err(Into::into)) - .boxed() - } + .boxed() }; Ok(stream) } Some(range) => { - // range is not none, the file maybe split into multiple parts to scan in parallel - // get footer_len firstly - let get_option = GetOptions { - range: Some(GetRange::Suffix(10)), - ..Default::default() - }; - let get_result = object_store - .get_opts(&partitioned_file.object_meta.location, get_option) - .await?; - let footer_len_buf = get_result.bytes().await?; - let footer_len = arrow_ipc::reader::read_footer_length( - footer_len_buf[..].try_into().unwrap(), + let footer_buf = read_file_footer(reader.as_ref()).await?; + let footer = arrow_ipc::root_as_footer(footer_buf.as_ref()).map_err( + |err| { + exec_datafusion_err!("Unable to get root as footer: {err:?}") + }, )?; - // read footer according to footer_len - let get_option = GetOptions { - range: Some(GetRange::Suffix(10 + (footer_len as u64))), - ..Default::default() - }; - let get_result = object_store - .get_opts(&partitioned_file.object_meta.location, get_option) - .await?; - let footer_buf = get_result.bytes().await?; - let footer = arrow_ipc::root_as_footer( - footer_buf[..footer_len].try_into().unwrap(), - ) - .map_err(|err| { - exec_datafusion_err!("Unable to get root as footer: {err:?}") - })?; // build decoder according to footer & projection let schema = arrow_ipc::convert::fb_to_schema(footer.schema().unwrap()); @@ -195,9 +181,7 @@ impl FileOpener for ArrowFileOpener { block_offset..block_offset + block_len }) .collect_vec(); - let dict_results = object_store - .get_ranges(&partitioned_file.object_meta.location, &dict_ranges) - .await?; + let dict_results = reader.read_ranges(dict_ranges).await?; for (dict_block, dict_result) in footer.dictionaries().iter().flatten().zip(dict_results) { @@ -228,12 +212,8 @@ impl FileOpener for ArrowFileOpener { }) .collect_vec(); - let recordbatch_results = object_store - .get_ranges( - &partitioned_file.object_meta.location, - &recordbatch_ranges, - ) - .await?; + let recordbatch_results = + reader.read_ranges(recordbatch_ranges).await?; let stream = futures::stream::iter( recordbatches @@ -291,19 +271,22 @@ impl ArrowSource { impl FileSource for ArrowSource { fn create_file_opener( &self, - object_store: Arc, + object_store: Arc, _base_config: &FileScanConfig, _partition: usize, + access_context: FileAccessContext, ) -> Result> { let split_projection = self.projection.clone(); let opener: Arc = match self.format { ArrowFormat::File => Arc::new(ArrowFileOpener { object_store, + access_context, projection: Some(split_projection.file_indices.clone()), }), ArrowFormat::Stream => Arc::new(ArrowStreamFileOpener { object_store, + access_context, projection: Some(split_projection.file_indices.clone()), }), }; @@ -510,24 +493,26 @@ impl ArrowOpener { } pub fn new_file_opener( - object_store: Arc, + object_store: Arc, projection: Option>, ) -> Self { Self { inner: Arc::new(ArrowFileOpener { object_store, + access_context: FileAccessContext::new("scan"), projection, }), } } pub fn new_stream_file_opener( - object_store: Arc, + object_store: Arc, projection: Option>, ) -> Self { Self { inner: Arc::new(ArrowStreamFileOpener { object_store, + access_context: FileAccessContext::new("scan"), projection, }), } @@ -548,10 +533,11 @@ mod tests { use arrow_ipc::reader::{FileReader, StreamReader}; use bytes::Bytes; use datafusion_datasource::file_scan_config::FileScanConfigBuilder; - use datafusion_execution::object_store::ObjectStoreUrl; + use datafusion_storage::StorageUrl; use object_store::memory::InMemory; use super::*; + use object_store::ObjectStoreExt; #[tokio::test] async fn test_file_opener_without_ranges() -> Result<()> { @@ -565,8 +551,14 @@ mod tests { let object_store = Arc::new(InMemory::new()); let partitioned_file = PartitionedFile::new(filename, file_size); object_store - .put(&partitioned_file.object_meta.location, bytes.into()) - .await?; + .put( + &object_store::path::Path::from( + partitioned_file.object_meta.location.as_ref(), + ), + bytes.into(), + ) + .await + .unwrap(); let schema = match FileReader::try_new(File::open(path_str)?, None) { Ok(reader) => reader.schema(), @@ -580,12 +572,17 @@ mod tests { }; let scan_config = FileScanConfigBuilder::new( - ObjectStoreUrl::local_filesystem(), + StorageUrl::local_filesystem(), source.clone(), ) .build(); - let file_opener = source.create_file_opener(object_store, &scan_config, 0)?; + let file_opener = source.create_file_opener( + test_utils::storage::object_store(object_store), + &scan_config, + 0, + FileAccessContext::default(), + )?; let mut stream = file_opener.open(partitioned_file)?.await?; assert!(stream.next().await.is_some()); @@ -611,20 +608,29 @@ mod tests { (file_size - 1) as i64, ); object_store - .put(&partitioned_file.object_meta.location, bytes.into()) - .await?; + .put( + &object_store::path::Path::from( + partitioned_file.object_meta.location.as_ref(), + ), + bytes.into(), + ) + .await + .unwrap(); let schema = FileReader::try_new(File::open(path_str)?, None)?.schema(); let source = Arc::new(ArrowSource::new_file_source(schema)); - let scan_config = FileScanConfigBuilder::new( - ObjectStoreUrl::local_filesystem(), - source.clone(), - ) - .build(); + let scan_config = + FileScanConfigBuilder::new(StorageUrl::local_filesystem(), source.clone()) + .build(); - let file_opener = source.create_file_opener(object_store, &scan_config, 0)?; + let file_opener = source.create_file_opener( + test_utils::storage::object_store(object_store), + &scan_config, + 0, + FileAccessContext::default(), + )?; let mut stream = file_opener.open(partitioned_file)?.await?; assert!(stream.next().await.is_some()); @@ -649,20 +655,29 @@ mod tests { (file_size - 1) as i64, ); object_store - .put(&partitioned_file.object_meta.location, bytes.into()) - .await?; + .put( + &object_store::path::Path::from( + partitioned_file.object_meta.location.as_ref(), + ), + bytes.into(), + ) + .await + .unwrap(); let schema = StreamReader::try_new(File::open(path_str)?, None)?.schema(); let source = Arc::new(ArrowSource::new_stream_file_source(schema)); - let scan_config = FileScanConfigBuilder::new( - ObjectStoreUrl::local_filesystem(), - source.clone(), - ) - .build(); + let scan_config = + FileScanConfigBuilder::new(StorageUrl::local_filesystem(), source.clone()) + .build(); - let file_opener = source.create_file_opener(object_store, &scan_config, 0)?; + let file_opener = source.create_file_opener( + test_utils::storage::object_store(object_store), + &scan_config, + 0, + FileAccessContext::default(), + )?; let result = file_opener.open(partitioned_file); assert!(result.is_err()); @@ -676,7 +691,7 @@ mod tests { let source = ArrowSource::new_stream_file_source(schema); let config = FileScanConfigBuilder::new( - ObjectStoreUrl::local_filesystem(), + StorageUrl::local_filesystem(), Arc::new(source.clone()) as Arc, ) .build(); @@ -706,11 +721,18 @@ mod tests { let object_store = Arc::new(InMemory::new()); let partitioned_file = PartitionedFile::new(filename, file_size); object_store - .put(&partitioned_file.object_meta.location, bytes.into()) - .await?; + .put( + &object_store::path::Path::from( + partitioned_file.object_meta.location.as_ref(), + ), + bytes.into(), + ) + .await + .unwrap(); let opener = ArrowStreamFileOpener { - object_store, + object_store: test_utils::storage::object_store(object_store), + access_context: FileAccessContext::default(), projection: Some(vec![0]), // just the first column }; diff --git a/datafusion/datasource-avro/Cargo.toml b/datafusion/datasource-avro/Cargo.toml index 70b675d63f427..5fc0c1e63080c 100644 --- a/datafusion/datasource-avro/Cargo.toml +++ b/datafusion/datasource-avro/Cargo.toml @@ -44,16 +44,18 @@ arrow = { workspace = true } arrow-avro = { workspace = true } async-trait = { workspace = true } bytes = { workspace = true } -datafusion-common = { workspace = true, features = ["object_store"] } +datafusion-common = { workspace = true } datafusion-datasource = { workspace = true } datafusion-physical-expr-adapter = { workspace = true } datafusion-physical-plan = { workspace = true } datafusion-proto-models = { workspace = true, optional = true } datafusion-session = { workspace = true } +datafusion-storage.workspace = true futures = { workspace = true } -object_store = { workspace = true } [dev-dependencies] +object_store = { workspace = true } +test-utils = { path = "../../test-utils" } # Note: add additional linter rules in lib.rs. # Rust does not support workspace + new linter rules in subcrates yet diff --git a/datafusion/datasource-avro/src/file_format.rs b/datafusion/datasource-avro/src/file_format.rs index 9d9d3279c0504..67c554636712b 100644 --- a/datafusion/datasource-avro/src/file_format.rs +++ b/datafusion/datasource-avro/src/file_format.rs @@ -39,7 +39,8 @@ use datafusion_physical_plan::ExecutionPlan; use datafusion_session::Session; use async_trait::async_trait; -use object_store::{GetResultPayload, ObjectMeta, ObjectStore, ObjectStoreExt}; + +use datafusion_storage::{FileInfo, StorageBinding}; #[derive(Default)] /// Factory struct used to create [`AvroFormat`] @@ -107,21 +108,18 @@ impl FileFormat for AvroFormat { async fn infer_schema( &self, _state: &dyn Session, - store: &Arc, - objects: &[ObjectMeta], + store: &Arc, + objects: &[FileInfo], ) -> Result { let mut schemas = vec![]; for object in objects { - let r = store.as_ref().get(&object.location).await?; - let schema = match r.payload { - GetResultPayload::File(mut file, _) => { - read_avro_schema_from_reader(&mut file)? - } - GetResultPayload::Stream(_) => { - // TODO: Fetching entire file to get schema is potentially wasteful - let data = r.bytes().await?; - read_avro_schema_from_reader(&mut data.as_ref())? - } + let reader = store + .open(object, datafusion_storage::FileAccessContext::new("schema")) + .await?; + let schema = { + // TODO: Fetching entire file to get schema is potentially wasteful + let data = datafusion_storage::collect_bytes(reader.stream(None)).await?; + read_avro_schema_from_reader(&mut data.as_ref())? }; schemas.push(schema); } @@ -132,9 +130,9 @@ impl FileFormat for AvroFormat { async fn infer_stats( &self, _state: &dyn Session, - _store: &Arc, + _store: &Arc, table_schema: SchemaRef, - _object: &ObjectMeta, + _object: &FileInfo, ) -> Result { Ok(Statistics::new_unknown(&table_schema)) } diff --git a/datafusion/datasource-avro/src/source.rs b/datafusion/datasource-avro/src/source.rs index e42a97e80c721..d325ccf6b85fd 100644 --- a/datafusion/datasource-avro/src/source.rs +++ b/datafusion/datasource-avro/src/source.rs @@ -32,7 +32,7 @@ use datafusion_physical_expr_adapter::BatchAdapterFactory; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; use datafusion_physical_plan::projection::ProjectionExprs; -use object_store::ObjectStore; +use datafusion_storage::StorageBinding; /// AvroSource holds the extra configuration that is necessary for opening avro files #[derive(Clone)] @@ -115,13 +115,15 @@ impl AvroSource { impl FileSource for AvroSource { fn create_file_opener( &self, - object_store: Arc, + object_store: Arc, _base_config: &FileScanConfig, _partition: usize, + access_context: datafusion_storage::FileAccessContext, ) -> Result> { let mut opener = Arc::new(private::AvroOpener { config: Arc::new(self.clone()), object_store, + access_context, }) as Arc; opener = ProjectionOpener::try_new( self.projection.clone(), @@ -249,73 +251,51 @@ impl AvroSource { mod private { use super::*; use std::io::BufReader; - use std::io::Seek; use bytes::Buf; use datafusion_datasource::{PartitionedFile, file_stream::FileOpenFuture}; use futures::StreamExt; - use object_store::{GetResultPayload, ObjectStore, ObjectStoreExt}; + + use datafusion_storage::StorageBinding; pub struct AvroOpener { pub config: Arc, - pub object_store: Arc, + pub object_store: Arc, + pub access_context: datafusion_storage::FileAccessContext, } impl FileOpener for AvroOpener { fn open(&self, partitioned_file: PartitionedFile) -> Result { let object_store = Arc::clone(&self.object_store); + let access_context = self.access_context.clone(); let config = Arc::clone(&self.config); let projected_file_schema = config.projected_file_schema(); Ok(Box::pin(async move { - let r = object_store - .get(&partitioned_file.object_meta.location) + let reader = object_store + .open(&partitioned_file.object_meta, access_context) .await?; - match r.payload { - GetResultPayload::File(mut file, _) => { - // Probe the writer schema first so logical projected columns can be - // translated to the writer-schema ordinals expected by `arrow-avro`. - let probe_reader = - config.open(BufReader::new(file.try_clone()?), None)?; - let writer_projection = config.writer_projection_for_schema( - probe_reader.schema().as_ref(), - projected_file_schema.as_ref(), - ); - file.rewind()?; - let reader = - config.open(BufReader::new(file), writer_projection)?; - let batch_adapter = - BatchAdapterFactory::new(Arc::clone(&projected_file_schema)) - .make_adapter(&reader.schema())?; - Ok(futures::stream::iter(reader) - .map(move |r| { - r.map_err(Into::into) - .and_then(|batch| batch_adapter.adapt_batch(&batch)) - }) - .boxed()) - } - GetResultPayload::Stream(_) => { - let bytes = r.bytes().await?; - // As above, inspect the writer schema before constructing the real - // reader so `with_projection` can use writer-schema ordinals. - let probe_reader = - config.open(BufReader::new(bytes.clone().reader()), None)?; - let writer_projection = config.writer_projection_for_schema( - probe_reader.schema().as_ref(), - projected_file_schema.as_ref(), - ); - let reader = config - .open(BufReader::new(bytes.reader()), writer_projection)?; - let batch_adapter = - BatchAdapterFactory::new(Arc::clone(&projected_file_schema)) - .make_adapter(&reader.schema())?; - Ok(futures::stream::iter(reader) - .map(move |r| { - r.map_err(Into::into) - .and_then(|batch| batch_adapter.adapt_batch(&batch)) - }) - .boxed()) - } + { + let bytes = + datafusion_storage::collect_bytes(reader.stream(None)).await?; + // Inspect the writer schema so projection uses writer-schema ordinals. + let probe_reader = + config.open(BufReader::new(bytes.clone().reader()), None)?; + let writer_projection = config.writer_projection_for_schema( + probe_reader.schema().as_ref(), + projected_file_schema.as_ref(), + ); + let reader = + config.open(BufReader::new(bytes.reader()), writer_projection)?; + let batch_adapter = + BatchAdapterFactory::new(Arc::clone(&projected_file_schema)) + .make_adapter(&reader.schema())?; + Ok(futures::stream::iter(reader) + .map(move |r| { + r.map_err(Into::into) + .and_then(|batch| batch_adapter.adapt_batch(&batch)) + }) + .boxed()) } })) } diff --git a/datafusion/datasource-csv/Cargo.toml b/datafusion/datasource-csv/Cargo.toml index 7e7195dfda9d5..c1f5013471675 100644 --- a/datafusion/datasource-csv/Cargo.toml +++ b/datafusion/datasource-csv/Cargo.toml @@ -42,7 +42,7 @@ proto = [ arrow = { workspace = true } async-trait = { workspace = true } bytes = { workspace = true } -datafusion-common = { workspace = true, features = ["object_store"] } +datafusion-common = { workspace = true } datafusion-common-runtime = { workspace = true } datafusion-datasource = { workspace = true } datafusion-execution = { workspace = true } @@ -51,8 +51,8 @@ datafusion-physical-expr-common = { workspace = true } datafusion-physical-plan = { workspace = true } datafusion-proto-models = { workspace = true, optional = true } datafusion-session = { workspace = true } +datafusion-storage.workspace = true futures = { workspace = true } -object_store = { workspace = true } regex = { workspace = true } tokio = { workspace = true } @@ -65,3 +65,7 @@ workspace = true [lib] name = "datafusion_datasource_csv" path = "src/mod.rs" + +[dev-dependencies] +object_store = { workspace = true } +test-utils = { path = "../../test-utils" } diff --git a/datafusion/datasource-csv/src/delimiter.rs b/datafusion/datasource-csv/src/delimiter.rs new file mode 100644 index 0000000000000..12b55956c096d --- /dev/null +++ b/datafusion/datasource-csv/src/delimiter.rs @@ -0,0 +1,257 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Groups CSV input into complete records for schema inference. + +use std::collections::VecDeque; + +use bytes::Bytes; +use futures::{Stream, StreamExt}; + +use datafusion_common::{Result, exec_datafusion_err}; + +/// The ASCII encoding of `"` +const QUOTE: u8 = b'"'; + +/// The ASCII encoding of `\n` +const NEWLINE: u8 = b'\n'; + +/// The ASCII encoding of `\` +const ESCAPE: u8 = b'\\'; + +/// [`LineDelimiter`] is provided with a stream of [`Bytes`] and returns an iterator +/// of [`Bytes`] containing a whole number of new line delimited records +#[derive(Debug, Default)] +struct LineDelimiter { + /// Complete chunks of [`Bytes`] + complete: VecDeque, + /// Remainder bytes that form the next record + remainder: Vec, + /// True if the last character was the escape character + is_escape: bool, + /// True if currently processing a quoted string + is_quote: bool, +} + +impl LineDelimiter { + /// Creates a new [`LineDelimiter`] with the provided delimiter + fn new() -> Self { + Self::default() + } + + /// Adds the next set of [`Bytes`] + fn push(&mut self, val: impl Into) { + let val: Bytes = val.into(); + + let is_escape = &mut self.is_escape; + let is_quote = &mut self.is_quote; + let mut record_ends = val.iter().enumerate().filter_map(|(idx, v)| { + if *is_escape { + *is_escape = false; + None + } else if *v == ESCAPE { + *is_escape = true; + None + } else if *v == QUOTE { + *is_quote = !*is_quote; + None + } else if *is_quote { + None + } else { + (*v == NEWLINE).then_some(idx + 1) + } + }); + + let start_offset = match self.remainder.is_empty() { + true => 0, + false => match record_ends.next() { + Some(idx) => { + self.remainder.extend_from_slice(&val[0..idx]); + self.complete + .push_back(Bytes::from(std::mem::take(&mut self.remainder))); + idx + } + None => { + self.remainder.extend_from_slice(&val); + return; + } + }, + }; + let end_offset = record_ends.next_back().unwrap_or(start_offset); + if start_offset != end_offset { + self.complete.push_back(val.slice(start_offset..end_offset)); + } + + if end_offset != val.len() { + self.remainder.extend_from_slice(&val[end_offset..]) + } + } + + /// Marks the end of the stream, delimiting any remaining bytes + /// + /// Returns `true` if there is no remaining data to be read + fn finish(&mut self) -> Result { + if !self.remainder.is_empty() { + if self.is_quote { + return Err(exec_datafusion_err!("encountered unterminated string")); + } + if self.is_escape { + return Err(exec_datafusion_err!( + "encountered trailing escape character" + )); + } + + self.complete + .push_back(Bytes::from(std::mem::take(&mut self.remainder))) + } + Ok(self.complete.is_empty()) + } +} + +impl Iterator for LineDelimiter { + type Item = Bytes; + + fn next(&mut self) -> Option { + self.complete.pop_front() + } +} + +/// Given a [`Stream`] of [`Bytes`] returns a [`Stream`] where each +/// yielded [`Bytes`] contains a whole number of new line delimited records +/// accounting for `\` style escapes and `"` quotes +pub fn newline_delimited_stream(s: S) -> impl Stream> +where + S: Stream> + Unpin, +{ + let delimiter = LineDelimiter::new(); + + futures::stream::unfold( + (s, delimiter, false), + |(mut s, mut delimiter, mut exhausted)| async move { + loop { + if let Some(next) = delimiter.next() { + return Some((Ok(next), (s, delimiter, exhausted))); + } else if exhausted { + return None; + } + + match s.next().await { + Some(Ok(bytes)) => delimiter.push(bytes), + Some(Err(e)) => return Some((Err(e), (s, delimiter, exhausted))), + None => { + exhausted = true; + match delimiter.finish() { + Ok(true) => return None, + Ok(false) => {} + Err(e) => return Some((Err(e), (s, delimiter, exhausted))), + } + } + } + } + }, + ) +} + +#[cfg(test)] +mod tests { + use futures::stream::{BoxStream, TryStreamExt}; + + use super::*; + + #[test] + fn test_delimiter() { + let mut delimiter = LineDelimiter::new(); + delimiter.push("hello\nworld"); + delimiter.push("\n\n"); + + assert_eq!(delimiter.next().unwrap(), Bytes::from("hello\n")); + assert_eq!(delimiter.next().unwrap(), Bytes::from("world\n")); + assert_eq!(delimiter.next().unwrap(), Bytes::from("\n")); + assert!(delimiter.next().is_none()); + } + + #[test] + fn test_delimiter_escaped() { + let mut delimiter = LineDelimiter::new(); + delimiter.push(""); + delimiter.push("fo\\\n\"foo"); + delimiter.push("bo\n\"bar\n"); + delimiter.push("\"he"); + delimiter.push("llo\"\n"); + assert_eq!( + delimiter.next().unwrap(), + Bytes::from("fo\\\n\"foobo\n\"bar\n") + ); + assert_eq!(delimiter.next().unwrap(), Bytes::from("\"hello\"\n")); + assert!(delimiter.next().is_none()); + + // Verify can push further data + delimiter.push("\"foo\nbar\",\"fiz\\\"inner\\\"\"\nhello"); + assert!(!delimiter.finish().unwrap()); + + assert_eq!( + delimiter.next().unwrap(), + Bytes::from("\"foo\nbar\",\"fiz\\\"inner\\\"\"\n") + ); + assert_eq!(delimiter.next().unwrap(), Bytes::from("hello")); + assert!(delimiter.finish().unwrap()); + assert!(delimiter.next().is_none()); + } + + #[tokio::test] + async fn test_delimiter_stream() { + let input = vec!["hello\nworld\nbin", "go\ncup", "cakes"]; + let input_stream = + futures::stream::iter(input.into_iter().map(|s| Ok(Bytes::from(s)))); + let stream = newline_delimited_stream(input_stream); + + let results: Vec<_> = stream.try_collect().await.unwrap(); + assert_eq!( + results, + vec![ + Bytes::from("hello\nworld\n"), + Bytes::from("bingo\n"), + Bytes::from("cupcakes") + ] + ) + } + #[tokio::test] + async fn test_delimiter_unfold_stream() { + let input_stream: BoxStream<'static, Result> = futures::stream::unfold( + VecDeque::from(["hello\nworld\nbin", "go\ncup", "cakes"]), + |mut input| async move { + if !input.is_empty() { + Some((Ok(Bytes::from(input.pop_front().unwrap())), input)) + } else { + None + } + }, + ) + .boxed(); + let stream = newline_delimited_stream(input_stream); + + let results: Vec<_> = stream.try_collect().await.unwrap(); + assert_eq!( + results, + vec![ + Bytes::from("hello\nworld\n"), + Bytes::from("bingo\n"), + Bytes::from("cupcakes") + ] + ) + } +} diff --git a/datafusion/datasource-csv/src/file_format.rs b/datafusion/datasource-csv/src/file_format.rs index a8c27369cb9c1..5f7309cb558b6 100644 --- a/datafusion/datasource-csv/src/file_format.rs +++ b/datafusion/datasource-csv/src/file_format.rs @@ -55,14 +55,13 @@ use datafusion_physical_expr_common::sort_expr::LexRequirement; use datafusion_physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan}; use datafusion_session::Session; +use crate::delimiter::newline_delimited_stream; use async_trait::async_trait; use bytes::{Buf, Bytes}; use datafusion_datasource::source::DataSourceExec; +use datafusion_storage::{FileInfo, StorageBinding}; use futures::stream::BoxStream; use futures::{Stream, StreamExt, TryStreamExt, pin_mut}; -use object_store::{ - ObjectMeta, ObjectStore, ObjectStoreExt, delimited::newline_delimited_stream, -}; use regex::Regex; #[derive(Default)] @@ -143,20 +142,20 @@ impl CsvFormat { /// Each returned `Bytes` has a whole number of newline delimited rows async fn read_to_delimited_chunks<'a>( &self, - store: &Arc, - object: &ObjectMeta, + store: &Arc, + object: &FileInfo, ) -> BoxStream<'a, Result> { // stream to only read as many rows as needed into memory let stream = store - .get(&object.location) + .open(object, datafusion_storage::FileAccessContext::new("schema")) .await - .map_err(|e| DataFusionError::ObjectStore(Box::new(e))); + .map_err(|e| DataFusionError::Storage(Box::new(e))); let stream = match stream { Ok(stream) => self .read_to_delimited_chunks_from_stream( stream - .into_stream() - .map_err(|e| DataFusionError::ObjectStore(Box::new(e))) + .stream(None) + .map_err(|e| DataFusionError::Storage(Box::new(e))) .boxed(), ) .map_err(DataFusionError::from) @@ -177,17 +176,7 @@ impl CsvFormat { let file_compression_type: FileCompressionType = self.options.compression.into(); let decoder = file_compression_type.convert_stream(stream); let stream = match decoder { - Ok(decoded_stream) => { - newline_delimited_stream(decoded_stream.map_err(|e| match e { - DataFusionError::ObjectStore(e) => *e, - err => object_store::Error::Generic { - store: "read to delimited chunks failed", - source: Box::new(err), - }, - })) - .map_err(DataFusionError::from) - .left_stream() - } + Ok(decoded_stream) => newline_delimited_stream(decoded_stream).left_stream(), Err(e) => { futures::stream::once(futures::future::ready(Err(e))).right_stream() } @@ -376,8 +365,8 @@ impl FileFormat for CsvFormat { async fn infer_schema( &self, state: &dyn Session, - store: &Arc, - objects: &[ObjectMeta], + store: &Arc, + objects: &[FileInfo], ) -> Result { let mut schemas = vec![]; @@ -423,9 +412,9 @@ impl FileFormat for CsvFormat { async fn infer_stats( &self, _state: &dyn Session, - _store: &Arc, + _store: &Arc, table_schema: SchemaRef, - _object: &ObjectMeta, + _object: &FileInfo, ) -> Result { Ok(Statistics::new_unknown(&table_schema)) } @@ -803,7 +792,7 @@ impl FileSink for CsvSink { context: &Arc, demux_task: SpawnedTask>, file_stream_rx: DemuxedStreamReceiver, - object_store: Arc, + object_store: Arc, ) -> Result { let builder = self.writer_options.writer_options.clone(); let header = builder.header(); @@ -924,7 +913,12 @@ impl CsvSink { "CsvSinkExecNode is missing required field 'sink'" ) })?; - let data_sink = CsvSink::try_from(proto_sink)?; + let mut data_sink = CsvSink::try_from(proto_sink)?; + data_sink.config.storage = Some( + ctx.task_ctx() + .runtime_env() + .storage(&data_sink.config.object_store_url)?, + ); let sort_order = DataSinkExec::decode_sort_order( sink_node.sort_order.as_ref(), ctx, diff --git a/datafusion/datasource-csv/src/mod.rs b/datafusion/datasource-csv/src/mod.rs index fdfee05d86a79..1d15d98173c91 100644 --- a/datafusion/datasource-csv/src/mod.rs +++ b/datafusion/datasource-csv/src/mod.rs @@ -29,7 +29,7 @@ use datafusion_common::Result; use datafusion_datasource::file_groups::FileGroup; use datafusion_datasource::file_scan_config::FileScanConfigBuilder; use datafusion_datasource::{file::FileSource, file_scan_config::FileScanConfig}; -use datafusion_execution::object_store::ObjectStoreUrl; +use datafusion_storage::StorageUrl; pub use file_format::*; /// Returns a [`FileScanConfig`] for given `file_groups` @@ -38,8 +38,10 @@ pub fn partitioned_csv_config( file_source: Arc, ) -> Result { Ok( - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + FileScanConfigBuilder::new(StorageUrl::local_filesystem(), file_source) .with_file_groups(file_groups) .build(), ) } + +mod delimiter; diff --git a/datafusion/datasource-csv/src/source.rs b/datafusion/datasource-csv/src/source.rs index 0bfa499d7a111..9c487be51daa6 100644 --- a/datafusion/datasource-csv/src/source.rs +++ b/datafusion/datasource-csv/src/source.rs @@ -21,7 +21,6 @@ use datafusion_datasource::boundary_stream::AlignedBoundaryStream; use datafusion_datasource::projection::{ProjectionOpener, SplitProjection}; use datafusion_physical_plan::projection::ProjectionExprs; use std::fmt; -use std::io::Read; use std::sync::Arc; use datafusion_datasource::decoder::{DecoderDeserializer, deserialize_stream}; @@ -39,15 +38,16 @@ use datafusion_common_runtime::JoinSet; use datafusion_datasource::file::FileSource; use datafusion_datasource::file_scan_config::FileScanConfig; use datafusion_execution::TaskContext; -use datafusion_physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet}; +use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; use datafusion_physical_plan::{ DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, }; use crate::file_format::CsvDecoder; +use datafusion_storage::FileAccessContext; +use datafusion_storage::StorageBinding; +use datafusion_storage::WriterOptions; use futures::{StreamExt, TryStreamExt}; -use object_store::buffered::BufWriter; -use object_store::{GetOptions, GetResultPayload, ObjectStore}; use tokio::io::AsyncWriteExt; /// A Config for [`CsvOpener`] @@ -59,11 +59,11 @@ use tokio::io::AsyncWriteExt; /// # use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; /// # use datafusion_datasource::PartitionedFile; /// # use datafusion_datasource_csv::source::CsvSource; -/// # use datafusion_execution::object_store::ObjectStoreUrl; +/// # use datafusion_storage::StorageUrl; /// # use datafusion_datasource::source::DataSourceExec; /// # use datafusion_common::config::CsvOptions; /// -/// # let object_store_url = ObjectStoreUrl::local_filesystem(); +/// # let object_store_url = StorageUrl::local_filesystem(); /// # let file_schema = Arc::new(Schema::empty()); /// /// let options = CsvOptions { @@ -180,10 +180,6 @@ impl CsvSource { } impl CsvSource { - fn open(&self, reader: R) -> Result> { - Ok(self.builder().build(reader)?) - } - fn builder(&self) -> csv::ReaderBuilder { let mut builder = csv::ReaderBuilder::new(Arc::clone(self.table_schema.file_schema())) @@ -214,8 +210,8 @@ impl CsvSource { pub struct CsvOpener { config: Arc, file_compression_type: FileCompressionType, - object_store: Arc, - partition_index: usize, + object_store: Arc, + access_context: FileAccessContext, } impl CsvOpener { @@ -223,13 +219,13 @@ impl CsvOpener { pub fn new( config: Arc, file_compression_type: FileCompressionType, - object_store: Arc, + object_store: Arc, ) -> Self { Self { config, file_compression_type, object_store, - partition_index: 0, + access_context: FileAccessContext::new("scan"), } } } @@ -243,15 +239,16 @@ impl From for Arc { impl FileSource for CsvSource { fn create_file_opener( &self, - object_store: Arc, + object_store: Arc, base_config: &FileScanConfig, - partition_index: usize, + _partition_index: usize, + access_context: FileAccessContext, ) -> Result> { let mut opener = Arc::new(CsvOpener { config: Arc::new(self.clone()), file_compression_type: base_config.file_compression_type, object_store, - partition_index, + access_context, }) as Arc; opener = ProjectionOpener::try_new( self.projection.clone(), @@ -413,15 +410,15 @@ impl FileOpener for CsvOpener { } let store = Arc::clone(&self.object_store); + let access_context = self.access_context.clone(); let terminator = self.config.terminator(); - let baseline_metrics = - BaselineMetrics::new(&self.config.metrics, self.partition_index); - Ok(Box::pin(async move { // Current partition contains bytes [start_byte, end_byte) (might contain incomplete lines at boundaries) let file_size = partitioned_file.object_meta.size; - let location = partitioned_file.object_meta.location; + let reader = store + .open(&partitioned_file.object_meta, access_context) + .await?; if let Some(file_range) = partitioned_file.range.as_ref() { let raw_start: u64 = file_range.start.try_into().map_err(|_| { @@ -438,14 +435,12 @@ impl FileOpener for CsvOpener { })?; let aligned_stream = AlignedBoundaryStream::new( - Arc::clone(&store), - location.clone(), + Arc::clone(&reader), raw_start, raw_end, file_size, terminator.unwrap_or(b'\n'), ) - .await? .map_err(DataFusionError::from); let decoder = config.builder().build_decoder(); @@ -459,39 +454,18 @@ impl FileOpener for CsvOpener { return Ok(stream.map_err(Into::into).boxed()); } - // No range specified — read the entire file - let options = GetOptions::default(); - let result = store.get_opts(&location, options).await?; - - match result.payload { - #[cfg(not(target_arch = "wasm32"))] - GetResultPayload::File(file, _) => { - let decoder = file_compression_type.convert_read(file)?; - let mut reader = config.open(decoder)?; - - // Use std::iter::from_fn to wrap execution of iterator's next() method. - let iterator = std::iter::from_fn(move || { - let mut timer = baseline_metrics.elapsed_compute().timer(); - let result = reader.next(); - timer.stop(); - result - }); - - Ok(futures::stream::iter(iterator) - .map(|r| r.map_err(Into::into)) - .boxed()) - } - GetResultPayload::Stream(s) => { - let decoder = config.builder().build_decoder(); - let s = s.map_err(DataFusionError::from); - let input = file_compression_type.convert_stream(s.boxed())?.fuse(); - - let stream = deserialize_stream( - input, - DecoderDeserializer::new(CsvDecoder::new(decoder)), - ); - Ok(stream.map_err(Into::into).boxed()) - } + let s = reader.stream(None); + + { + let decoder = config.builder().build_decoder(); + let s = s.map_err(DataFusionError::from); + let input = file_compression_type.convert_stream(s.boxed())?.fuse(); + + let stream = deserialize_stream( + input, + DecoderDeserializer::new(CsvDecoder::new(decoder)), + ); + Ok(stream.map_err(Into::into).boxed()) } })) } @@ -504,8 +478,8 @@ pub async fn plan_to_csv( ) -> Result<()> { let path = path.as_ref(); let parsed = ListingTableUrl::parse(path)?; - let object_store_url = parsed.object_store(); - let store = task_ctx.runtime_env().object_store(&object_store_url)?; + let object_store_url = parsed.storage_url(); + let store = task_ctx.runtime_env().storage(&object_store_url)?; let writer_buffer_size = task_ctx .session_config() .options() @@ -516,12 +490,19 @@ pub async fn plan_to_csv( let storeref = Arc::clone(&store); let plan: Arc = Arc::clone(&plan); let filename = format!("{}/part-{i}.csv", parsed.prefix()); - let file = object_store::path::Path::parse(filename)?; + let file = datafusion_storage::path::Path::parse(filename)?; let mut stream = plan.execute(i, Arc::clone(&task_ctx))?; join_set.spawn(async move { - let mut buf_writer = - BufWriter::with_capacity(storeref, file.clone(), writer_buffer_size); + let mut buf_writer = storeref + .writer( + &file, + WriterOptions { + buffer_size: Some(writer_buffer_size), + }, + FileAccessContext::new("write"), + ) + .await?; let mut buffer = Vec::with_capacity(1024); //only write headers on first iteration let mut write_headers = true; diff --git a/datafusion/datasource-json/Cargo.toml b/datafusion/datasource-json/Cargo.toml index 04192083f583a..cf53d480d7e49 100644 --- a/datafusion/datasource-json/Cargo.toml +++ b/datafusion/datasource-json/Cargo.toml @@ -42,7 +42,7 @@ proto = [ arrow = { workspace = true } async-trait = { workspace = true } bytes = { workspace = true } -datafusion-common = { workspace = true, features = ["object_store"] } +datafusion-common = { workspace = true } datafusion-common-runtime = { workspace = true } datafusion-datasource = { workspace = true } datafusion-execution = { workspace = true } @@ -51,8 +51,8 @@ datafusion-physical-expr-common = { workspace = true } datafusion-physical-plan = { workspace = true } datafusion-proto-models = { workspace = true, optional = true } datafusion-session = { workspace = true } +datafusion-storage.workspace = true futures = { workspace = true } -object_store = { workspace = true } tokio = { workspace = true } tokio-stream = { workspace = true, features = ["sync"] } @@ -63,7 +63,9 @@ tokio-stream = { workspace = true, features = ["sync"] } workspace = true [dev-dependencies] +object_store = { workspace = true } serde_json = { workspace = true } +test-utils = { path = "../../test-utils" } [lib] name = "datafusion_datasource_json" diff --git a/datafusion/datasource-json/src/file_format.rs b/datafusion/datasource-json/src/file_format.rs index aef89560a3b5f..85209a1234f09 100644 --- a/datafusion/datasource-json/src/file_format.rs +++ b/datafusion/datasource-json/src/file_format.rs @@ -60,7 +60,8 @@ use datafusion_session::Session; use crate::utils::JsonArrayToNdjsonReader; use async_trait::async_trait; -use object_store::{GetResultPayload, ObjectMeta, ObjectStore, ObjectStoreExt}; + +use datafusion_storage::{FileInfo, StorageBinding}; #[derive(Default)] /// Factory struct used to create [JsonFormat] @@ -255,8 +256,8 @@ impl FileFormat for JsonFormat { async fn infer_schema( &self, _state: &dyn Session, - store: &Arc, - objects: &[ObjectMeta], + store: &Arc, + objects: &[FileInfo], ) -> Result { let mut schemas = Vec::new(); let mut records_to_read = self @@ -272,53 +273,30 @@ impl FileFormat for JsonFormat { break; } - let r = store.as_ref().get(&object.location).await?; - - let (schema, records_consumed) = match r.payload { - #[cfg(not(target_arch = "wasm32"))] - GetResultPayload::File(file, _) => { - let decoder = file_compression_type.convert_read(file)?; - let reader = BufReader::new(decoder); - - if newline_delimited { - // NDJSON: use ValueIter directly - let iter = ValueIter::new(reader, None); - let mut count = 0; - let schema = - infer_json_schema_from_iterator(iter.take_while(|_| { - let should_take = count < records_to_read; - if should_take { - count += 1; - } - should_take - }))?; - (schema, count) - } else { - // JSON array format: use streaming converter - infer_schema_from_json_array(reader, records_to_read)? - } - } - GetResultPayload::Stream(_) => { - let data = r.bytes().await?; - let decoder = file_compression_type.convert_read(data.reader())?; - let reader = BufReader::new(decoder); - - if newline_delimited { - let iter = ValueIter::new(reader, None); - let mut count = 0; - let schema = - infer_json_schema_from_iterator(iter.take_while(|_| { - let should_take = count < records_to_read; - if should_take { - count += 1; - } - should_take - }))?; - (schema, count) - } else { - // JSON array format: use streaming converter - infer_schema_from_json_array(reader, records_to_read)? - } + let reader = store + .open(object, datafusion_storage::FileAccessContext::new("schema")) + .await?; + + let (schema, records_consumed) = { + let data = datafusion_storage::collect_bytes(reader.stream(None)).await?; + let decoder = file_compression_type.convert_read(data.reader())?; + let reader = BufReader::new(decoder); + + if newline_delimited { + let iter = ValueIter::new(reader, None); + let mut count = 0; + let schema = + infer_json_schema_from_iterator(iter.take_while(|_| { + let should_take = count < records_to_read; + if should_take { + count += 1; + } + should_take + }))?; + (schema, count) + } else { + // JSON array format: use streaming converter + infer_schema_from_json_array(reader, records_to_read)? } }; @@ -334,9 +312,9 @@ impl FileFormat for JsonFormat { async fn infer_stats( &self, _state: &dyn Session, - _store: &Arc, + _store: &Arc, table_schema: SchemaRef, - _object: &ObjectMeta, + _object: &FileInfo, ) -> Result { Ok(Statistics::new_unknown(&table_schema)) } @@ -461,7 +439,7 @@ impl FileSink for JsonSink { context: &Arc, demux_task: SpawnedTask>, file_stream_rx: DemuxedStreamReceiver, - object_store: Arc, + object_store: Arc, ) -> Result { let serializer = Arc::new(JsonSerializer::new()) as _; spawn_writer_tasks_and_join( @@ -598,7 +576,12 @@ impl JsonSink { "JsonSinkExecNode is missing required field 'sink'" ) })?; - let data_sink = JsonSink::try_from(proto_sink)?; + let mut data_sink = JsonSink::try_from(proto_sink)?; + data_sink.config.storage = Some( + ctx.task_ctx() + .runtime_env() + .storage(&data_sink.config.object_store_url)?, + ); let sort_order = DataSinkExec::decode_sort_order( sort_order.as_ref(), ctx, diff --git a/datafusion/datasource-json/src/source.rs b/datafusion/datasource-json/src/source.rs index f4c7021c86bdd..e75c0a7230adf 100644 --- a/datafusion/datasource-json/src/source.rs +++ b/datafusion/datasource-json/src/source.rs @@ -17,7 +17,6 @@ //! Execution plan for reading JSON files (line-delimited and array formats) -use std::io::BufReader; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; @@ -46,10 +45,12 @@ use datafusion_datasource::file_scan_config::FileScanConfig; use datafusion_execution::TaskContext; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; +use datafusion_storage::FileAccessContext; +use datafusion_storage::StorageBinding; +use datafusion_storage::WriterOptions; use futures::{Stream, StreamExt, TryStreamExt}; -use object_store::buffered::BufWriter; -use object_store::{GetOptions, GetResultPayload, ObjectStore}; use tokio::io::AsyncWriteExt; + use tokio_stream::wrappers::ReceiverStream; /// Channel buffer size for streaming JSON array processing. @@ -101,7 +102,8 @@ pub struct JsonOpener { batch_size: usize, projected_schema: SchemaRef, file_compression_type: FileCompressionType, - object_store: Arc, + object_store: Arc, + access_context: FileAccessContext, /// When `true` (default), expects newline-delimited JSON (NDJSON). /// When `false`, expects JSON array format `[{...}, {...}]`. newline_delimited: bool, @@ -113,7 +115,7 @@ impl JsonOpener { batch_size: usize, projected_schema: SchemaRef, file_compression_type: FileCompressionType, - object_store: Arc, + object_store: Arc, newline_delimited: bool, ) -> Self { Self { @@ -121,6 +123,7 @@ impl JsonOpener { projected_schema, file_compression_type, object_store, + access_context: FileAccessContext::new("scan"), newline_delimited, } } @@ -175,9 +178,10 @@ impl From for Arc { impl FileSource for JsonSource { fn create_file_opener( &self, - object_store: Arc, + object_store: Arc, base_config: &FileScanConfig, _partition: usize, + access_context: FileAccessContext, ) -> Result> { // Get the projected file schema for JsonOpener let file_schema = self.table_schema.file_schema(); @@ -191,6 +195,7 @@ impl FileSource for JsonSource { projected_schema, file_compression_type: base_config.file_compression_type, object_store, + access_context, newline_delimited: self.newline_delimited, }) as Arc; @@ -345,6 +350,7 @@ impl FileOpener for JsonOpener { /// Note: JSON array format does not support range-based scanning. fn open(&self, partitioned_file: PartitionedFile) -> Result { let store = Arc::clone(&self.object_store); + let access_context = self.access_context.clone(); let schema = Arc::clone(&self.projected_schema); let batch_size = self.batch_size; let file_compression_type = self.file_compression_type.to_owned(); @@ -361,7 +367,9 @@ impl FileOpener for JsonOpener { Ok(Box::pin(async move { let file_size = partitioned_file.object_meta.size; - let location = &partitioned_file.object_meta.location; + let reader = store + .open(&partitioned_file.object_meta, access_context) + .await?; if let Some(file_range) = partitioned_file.range.as_ref() { let raw_start: u64 = file_range.start.try_into().map_err(|_| { @@ -378,14 +386,12 @@ impl FileOpener for JsonOpener { })?; let aligned_stream = AlignedBoundaryStream::new( - Arc::clone(&store), - location.clone(), + Arc::clone(&reader), raw_start, raw_end, file_size, b'\n', ) - .await? .map_err(DataFusionError::from); let decoder = ReaderBuilder::new(schema) @@ -401,155 +407,118 @@ impl FileOpener for JsonOpener { return Ok(stream.map_err(Into::into).boxed()); } - // No range specified — read the entire file - let options = GetOptions::default(); - let result = store.get_opts(location, options).await?; - - match result.payload { - #[cfg(not(target_arch = "wasm32"))] - GetResultPayload::File(file, _) => { - let bytes = file_compression_type.convert_read(file)?; - - if newline_delimited { - // NDJSON: use BufReader directly - let reader = BufReader::new(bytes); - let arrow_reader = ReaderBuilder::new(schema) - .with_batch_size(batch_size) - .build(reader)?; - - Ok(futures::stream::iter(arrow_reader) - .map(|r| r.map_err(Into::into)) - .boxed()) - } else { - // JSON array format: wrap with streaming converter - let ndjson_reader = JsonArrayToNdjsonReader::with_capacity( - bytes, + let s = reader.stream(None); + + { + if newline_delimited { + // Newline-delimited JSON (NDJSON) streaming reader + let s = s.map_err(DataFusionError::from); + let decoder = ReaderBuilder::new(schema) + .with_batch_size(batch_size) + .build_decoder()?; + let input = file_compression_type.convert_stream(s.boxed())?.fuse(); + let stream = deserialize_stream( + input, + DecoderDeserializer::new(JsonDecoder::new(decoder)), + ); + Ok(stream.map_err(Into::into).boxed()) + } else { + // JSON array format: streaming conversion with channel-based byte transfer + // + // Architecture: + // 1. Async task reads from object store stream, decompresses, sends to channel + // 2. Blocking task receives bytes, converts JSON array to NDJSON, parses to Arrow + // 3. RecordBatches are sent back via another channel + // + // Memory budget (~32MB): + // - sync_channel: CHANNEL_BUFFER_SIZE chunks (~16MB) + // - JsonArrayToNdjsonReader: 2 × JSON_CONVERTER_BUFFER_SIZE (~4MB) + // - Arrow JsonReader internal buffer (~8MB) + // - Miscellaneous (~4MB) + + let s = s.map_err(DataFusionError::from); + let decompressed_stream = + file_compression_type.convert_stream(s.boxed())?; + + // Channel for bytes: async producer -> blocking consumer + // Uses tokio::sync::mpsc so the async send never blocks a + // tokio worker thread; the consumer calls blocking_recv() + // inside spawn_blocking. + let (byte_tx, byte_rx) = + tokio::sync::mpsc::channel::(CHANNEL_BUFFER_SIZE); + + // Channel for results: sync producer -> async consumer + let (result_tx, result_rx) = tokio::sync::mpsc::channel(2); + let error_tx = result_tx.clone(); + + // Async task: read from object store stream and send bytes to channel + // Store the SpawnedTask to keep it alive until stream is dropped + let read_task = SpawnedTask::spawn(async move { + tokio::pin!(decompressed_stream); + while let Some(chunk) = decompressed_stream.next().await { + match chunk { + Ok(bytes) => { + if byte_tx.send(bytes).await.is_err() { + break; // Consumer dropped + } + } + Err(e) => { + let _ = error_tx + .send(Err( + arrow::error::ArrowError::ExternalError( + Box::new(e), + ), + )) + .await; + break; + } + } + } + // byte_tx dropped here, signals EOF to ChannelReader + }); + + // Blocking task: receive bytes from channel and parse JSON + // Store the SpawnedTask to keep it alive until stream is dropped + let parse_task = SpawnedTask::spawn_blocking(move || { + let channel_reader = ChannelReader::new(byte_rx); + let mut ndjson_reader = JsonArrayToNdjsonReader::with_capacity( + channel_reader, JSON_CONVERTER_BUFFER_SIZE, ); - let arrow_reader = ReaderBuilder::new(schema) - .with_batch_size(batch_size) - .build(ndjson_reader)?; - Ok(futures::stream::iter(arrow_reader) - .map(|r| r.map_err(Into::into)) - .boxed()) - } - } - GetResultPayload::Stream(s) => { - if newline_delimited { - // Newline-delimited JSON (NDJSON) streaming reader - let s = s.map_err(DataFusionError::from); - let decoder = ReaderBuilder::new(schema) + match ReaderBuilder::new(schema) .with_batch_size(batch_size) - .build_decoder()?; - let input = - file_compression_type.convert_stream(s.boxed())?.fuse(); - let stream = deserialize_stream( - input, - DecoderDeserializer::new(JsonDecoder::new(decoder)), - ); - Ok(stream.map_err(Into::into).boxed()) - } else { - // JSON array format: streaming conversion with channel-based byte transfer - // - // Architecture: - // 1. Async task reads from object store stream, decompresses, sends to channel - // 2. Blocking task receives bytes, converts JSON array to NDJSON, parses to Arrow - // 3. RecordBatches are sent back via another channel - // - // Memory budget (~32MB): - // - sync_channel: CHANNEL_BUFFER_SIZE chunks (~16MB) - // - JsonArrayToNdjsonReader: 2 × JSON_CONVERTER_BUFFER_SIZE (~4MB) - // - Arrow JsonReader internal buffer (~8MB) - // - Miscellaneous (~4MB) - - let s = s.map_err(DataFusionError::from); - let decompressed_stream = - file_compression_type.convert_stream(s.boxed())?; - - // Channel for bytes: async producer -> blocking consumer - // Uses tokio::sync::mpsc so the async send never blocks a - // tokio worker thread; the consumer calls blocking_recv() - // inside spawn_blocking. - let (byte_tx, byte_rx) = tokio::sync::mpsc::channel::( - CHANNEL_BUFFER_SIZE, - ); - - // Channel for results: sync producer -> async consumer - let (result_tx, result_rx) = tokio::sync::mpsc::channel(2); - let error_tx = result_tx.clone(); - - // Async task: read from object store stream and send bytes to channel - // Store the SpawnedTask to keep it alive until stream is dropped - let read_task = SpawnedTask::spawn(async move { - tokio::pin!(decompressed_stream); - while let Some(chunk) = decompressed_stream.next().await { - match chunk { - Ok(bytes) => { - if byte_tx.send(bytes).await.is_err() { - break; // Consumer dropped - } - } - Err(e) => { - let _ = error_tx - .send(Err( - arrow::error::ArrowError::ExternalError( - Box::new(e), - ), - )) - .await; - break; + .build(&mut ndjson_reader) + { + Ok(arrow_reader) => { + for batch_result in arrow_reader { + if result_tx.blocking_send(batch_result).is_err() { + break; // Receiver dropped } } } - // byte_tx dropped here, signals EOF to ChannelReader - }); - - // Blocking task: receive bytes from channel and parse JSON - // Store the SpawnedTask to keep it alive until stream is dropped - let parse_task = SpawnedTask::spawn_blocking(move || { - let channel_reader = ChannelReader::new(byte_rx); - let mut ndjson_reader = - JsonArrayToNdjsonReader::with_capacity( - channel_reader, - JSON_CONVERTER_BUFFER_SIZE, - ); - - match ReaderBuilder::new(schema) - .with_batch_size(batch_size) - .build(&mut ndjson_reader) - { - Ok(arrow_reader) => { - for batch_result in arrow_reader { - if result_tx.blocking_send(batch_result).is_err() - { - break; // Receiver dropped - } - } - } - Err(e) => { - let _ = result_tx.blocking_send(Err(e)); - } - } - - // Validate the JSON array was properly formed - if let Err(e) = ndjson_reader.validate_complete() { - let _ = result_tx.blocking_send(Err( - arrow::error::ArrowError::JsonError(e.to_string()), - )); + Err(e) => { + let _ = result_tx.blocking_send(Err(e)); } - // result_tx dropped here, closes the stream - }); - - // Wrap in JsonArrayStream to keep tasks alive until stream is consumed - let stream = JsonArrayStream { - inner: ReceiverStream::new(result_rx), - _read_task: read_task, - _parse_task: parse_task, - }; - - Ok(stream.map(|r| r.map_err(Into::into)).boxed()) - } + } + + // Validate the JSON array was properly formed + if let Err(e) = ndjson_reader.validate_complete() { + let _ = result_tx.blocking_send(Err( + arrow::error::ArrowError::JsonError(e.to_string()), + )); + } + // result_tx dropped here, closes the stream + }); + + // Wrap in JsonArrayStream to keep tasks alive until stream is consumed + let stream = JsonArrayStream { + inner: ReceiverStream::new(result_rx), + _read_task: read_task, + _parse_task: parse_task, + }; + + Ok(stream.map(|r| r.map_err(Into::into)).boxed()) } } })) @@ -563,8 +532,8 @@ pub async fn plan_to_json( ) -> Result<()> { let path = path.as_ref(); let parsed = ListingTableUrl::parse(path)?; - let object_store_url = parsed.object_store(); - let store = task_ctx.runtime_env().object_store(&object_store_url)?; + let object_store_url = parsed.storage_url(); + let store = task_ctx.runtime_env().storage(&object_store_url)?; let writer_buffer_size = task_ctx .session_config() .options() @@ -575,12 +544,19 @@ pub async fn plan_to_json( let storeref = Arc::clone(&store); let plan: Arc = Arc::clone(&plan); let filename = format!("{}/part-{i}.json", parsed.prefix()); - let file = object_store::path::Path::parse(filename)?; + let file = datafusion_storage::path::Path::parse(filename)?; let mut stream = plan.execute(i, Arc::clone(&task_ctx))?; join_set.spawn(async move { - let mut buf_writer = - BufWriter::with_capacity(storeref, file.clone(), writer_buffer_size); + let mut buf_writer = storeref + .writer( + &file, + WriterOptions { + buffer_size: Some(writer_buffer_size), + }, + FileAccessContext::new("write"), + ) + .await?; let mut buffer = Vec::with_capacity(1024); while let Some(batch) = stream.next().await.transpose()? { @@ -621,6 +597,7 @@ mod tests { use arrow::record_batch::RecordBatch; use bytes::Bytes; use datafusion_datasource::FileRange; + use object_store::ObjectStore; use object_store::memory::InMemory; use object_store::path::Path; use object_store::{ObjectStoreExt, PutPayload}; @@ -643,17 +620,18 @@ mod tests { let path = Path::from("test.json"); store .put(&path, PutPayload::from_static(json_data.as_bytes())) - .await?; + .await + .unwrap(); let opener = JsonOpener::new( 1024, test_schema(), FileCompressionType::UNCOMPRESSED, - store.clone(), + test_utils::storage::object_store(store.clone()), false, // JSON array format ); - let meta = store.head(&path).await?; + let meta = store.head(&path).await.unwrap(); let file = PartitionedFile::new(path.to_string(), meta.size); let stream = opener.open(file)?.await?; @@ -675,17 +653,18 @@ mod tests { let path = Path::from("test_stream.json"); store .put(&path, PutPayload::from_static(json_data.as_bytes())) - .await?; + .await + .unwrap(); let opener = JsonOpener::new( 2, // small batch size to test multiple batches test_schema(), FileCompressionType::UNCOMPRESSED, - store.clone(), + test_utils::storage::object_store(store.clone()), false, // JSON array format ); - let meta = store.head(&path).await?; + let meta = store.head(&path).await.unwrap(); let file = PartitionedFile::new(path.to_string(), meta.size); let stream = opener.open(file)?.await?; @@ -714,17 +693,18 @@ mod tests { let path = Path::from("nested.json"); store .put(&path, PutPayload::from_static(json_data.as_bytes())) - .await?; + .await + .unwrap(); let opener = JsonOpener::new( 1024, schema, FileCompressionType::UNCOMPRESSED, - store.clone(), + test_utils::storage::object_store(store.clone()), false, ); - let meta = store.head(&path).await?; + let meta = store.head(&path).await.unwrap(); let file = PartitionedFile::new(path.to_string(), meta.size); let stream = opener.open(file)?.await?; @@ -744,17 +724,18 @@ mod tests { let path = Path::from("empty.json"); store .put(&path, PutPayload::from_static(json_data.as_bytes())) - .await?; + .await + .unwrap(); let opener = JsonOpener::new( 1024, test_schema(), FileCompressionType::UNCOMPRESSED, - store.clone(), + test_utils::storage::object_store(store.clone()), false, ); - let meta = store.head(&path).await?; + let meta = store.head(&path).await.unwrap(); let file = PartitionedFile::new(path.to_string(), meta.size); let stream = opener.open(file)?.await?; @@ -780,7 +761,7 @@ mod tests { 1024, test_schema(), FileCompressionType::UNCOMPRESSED, - store.clone(), + test_utils::storage::object_store(store.clone()), false, // JSON array format ); @@ -810,17 +791,18 @@ mod tests { let path = Path::from("test.ndjson"); store .put(&path, PutPayload::from_static(json_data.as_bytes())) - .await?; + .await + .unwrap(); let opener = JsonOpener::new( 1024, test_schema(), FileCompressionType::UNCOMPRESSED, - store.clone(), + test_utils::storage::object_store(store.clone()), true, // NDJSON format ); - let meta = store.head(&path).await?; + let meta = store.head(&path).await.unwrap(); let file = PartitionedFile::new(path.to_string(), meta.size); let stream = opener.open(file)?.await?; @@ -848,17 +830,18 @@ mod tests { let path = Path::from("large.json"); store .put(&path, PutPayload::from(Bytes::from(json_data))) - .await?; + .await + .unwrap(); let opener = JsonOpener::new( 100, // batch size of 100 test_schema(), FileCompressionType::UNCOMPRESSED, - store.clone(), + test_utils::storage::object_store(store.clone()), false, ); - let meta = store.head(&path).await?; + let meta = store.head(&path).await.unwrap(); let file = PartitionedFile::new(path.to_string(), meta.size); let stream = opener.open(file)?.await?; @@ -889,17 +872,18 @@ mod tests { let path = Path::from("cancel_test.json"); store .put(&path, PutPayload::from(Bytes::from(json_data))) - .await?; + .await + .unwrap(); let opener = JsonOpener::new( 10, // small batch size test_schema(), FileCompressionType::UNCOMPRESSED, - store.clone(), + test_utils::storage::object_store(store.clone()), false, ); - let meta = store.head(&path).await?; + let meta = store.head(&path).await.unwrap(); let file = PartitionedFile::new(path.to_string(), meta.size); let mut stream = opener.open(file)?.await?; @@ -935,7 +919,7 @@ mod tests { let start = (p as u64 * file_size) / num_partitions as u64; let end = ((p as u64 + 1) * file_size) / num_partitions as u64; - let meta = store.head(path).await?; + let meta = store.head(path).await.unwrap(); let mut file = PartitionedFile::new(path.to_string(), meta.size); file.range = Some(FileRange { start: start as i64, @@ -946,7 +930,7 @@ mod tests { 1024, test_schema(), FileCompressionType::UNCOMPRESSED, - Arc::clone(&store), + test_utils::storage::object_store(store.clone()), true, ); @@ -987,7 +971,7 @@ mod tests { for num_partitions in get_partition_splits() { let batches = collect_partitioned_batches( - Arc::clone(&store), + store.clone(), &path, file_size, num_partitions, @@ -1059,7 +1043,7 @@ mod tests { for num_partitions in get_partition_splits() { let batches = collect_partitioned_batches( - Arc::clone(&store), + store.clone(), &path, file_size, num_partitions, @@ -1115,7 +1099,7 @@ mod tests { for num_partitions in get_partition_splits() { let batches = collect_partitioned_batches( - Arc::clone(&store), + store.clone(), &path, file_size, num_partitions, diff --git a/datafusion/datasource-parquet/Cargo.toml b/datafusion/datasource-parquet/Cargo.toml index a2589af19a6ee..0fcb32342562b 100644 --- a/datafusion/datasource-parquet/Cargo.toml +++ b/datafusion/datasource-parquet/Cargo.toml @@ -35,7 +35,7 @@ arrow = { workspace = true } arrow-schema = { workspace = true } async-trait = { workspace = true } bytes = { workspace = true } -datafusion-common = { workspace = true, features = ["object_store", "parquet"] } +datafusion-common = { workspace = true, features = ["parquet"] } datafusion-common-runtime = { workspace = true } datafusion-datasource = { workspace = true } datafusion-execution = { workspace = true } @@ -49,20 +49,23 @@ datafusion-physical-plan = { workspace = true } datafusion-proto-models = { workspace = true, optional = true } datafusion-pruning = { workspace = true } datafusion-session = { workspace = true } +datafusion-storage = { workspace = true } futures = { workspace = true } itertools = { workspace = true } log = { workspace = true } -object_store = { workspace = true } parking_lot = { workspace = true } parquet = { workspace = true } tokio = { workspace = true } +url = { workspace = true } [dev-dependencies] chrono = { workspace = true } criterion = { workspace = true } datafusion-functions = { workspace = true } datafusion-functions-nested = { workspace = true } +object_store = { workspace = true } tempfile = { workspace = true } +test-utils = { path = "../../test-utils" } # Note: add additional linter rules in lib.rs. # Rust does not support workspace + new linter rules in subcrates yet @@ -75,7 +78,9 @@ name = "datafusion_datasource_parquet" path = "src/mod.rs" [features] +default = [] proto = [ + "dep:datafusion-proto-models", "datafusion-datasource/proto", "datafusion-physical-plan/proto", diff --git a/datafusion/datasource-parquet/src/bloom_filter.rs b/datafusion/datasource-parquet/src/bloom_filter.rs index 2384f95df9885..f2bec35a9bd33 100644 --- a/datafusion/datasource-parquet/src/bloom_filter.rs +++ b/datafusion/datasource-parquet/src/bloom_filter.rs @@ -655,10 +655,14 @@ mod tests { let file_metrics = ParquetFileMetrics::new(0, object_meta.location.as_ref(), &metrics); let store: Arc = Arc::new(in_memory); - let partitioned_file = PartitionedFile::new_from_meta(object_meta); + let partitioned_file = + PartitionedFile::new_from_meta(test_utils::storage::file_info(object_meta)); - let reader = - ParquetFileReader::new(file_metrics.clone(), store, partitioned_file); + let reader = ParquetFileReader::new( + file_metrics.clone(), + test_utils::storage::object_store(store), + partitioned_file, + ); let mut builder = ParquetRecordBatchStreamBuilder::new(reader).await.unwrap(); let access_plan = ParquetAccessPlan::new_all(builder.metadata().num_row_groups()); diff --git a/datafusion/datasource-parquet/src/file_format.rs b/datafusion/datasource-parquet/src/file_format.rs index 1b25a2c632510..0be02f47d642f 100644 --- a/datafusion/datasource-parquet/src/file_format.rs +++ b/datafusion/datasource-parquet/src/file_format.rs @@ -20,7 +20,6 @@ use std::collections::HashSet; use std::fmt; use std::fmt::Debug; -use std::ops::Range; use std::sync::Arc; // Re-export so the historical `file_format::*` paths still resolve. @@ -60,20 +59,15 @@ use datafusion_physical_plan::ExecutionPlan; use datafusion_session::Session; use crate::metadata::{DFParquetMetadata, lex_ordering_to_sorting_columns}; -use crate::reader::CachedParquetFileReaderFactory; use crate::source::{ ParquetSource, parse_coerce_int96_string, parse_coerce_int96_tz_string, }; use async_trait::async_trait; -use bytes::Bytes; use datafusion_datasource::source::DataSourceExec; use datafusion_execution::cache::cache_manager::FileMetadataCache; -use futures::future::BoxFuture; -use futures::{FutureExt, StreamExt, TryStreamExt}; -use object_store::path::Path; -use object_store::{ObjectMeta, ObjectStore, ObjectStoreExt}; -use parquet::arrow::async_reader::MetadataFetch; -use parquet::errors::ParquetError; +use datafusion_storage::path::Path; +use datafusion_storage::{FileAccessContext, FileInfo, StorageBinding}; +use futures::{StreamExt, TryStreamExt}; use parquet::file::metadata::ParquetMetaData; #[derive(Default)] @@ -334,8 +328,8 @@ impl FileFormat for ParquetFormat { async fn infer_schema( &self, state: &dyn Session, - store: &Arc, - objects: &[ObjectMeta], + store: &Arc, + objects: &[FileInfo], ) -> Result { let coerce_int96 = match self.coerce_int96() { Some(time_unit) => Some(parse_coerce_int96_string(time_unit.as_str())?), @@ -360,10 +354,16 @@ impl FileFormat for ParquetFormat { &object.location, ) .await?; - let result = DFParquetMetadata::new(store.as_ref(), object) + let reader = store + .open(object, FileAccessContext::new("parquet-schema")) + .await?; + let result = DFParquetMetadata::new(reader.as_ref(), object) .with_metadata_size_hint(self.metadata_size_hint()) .with_decryption_properties(file_decryption_properties) - .with_file_metadata_cache(Some(Arc::clone(&file_metadata_cache))) + .with_file_metadata_cache( + Some(Arc::clone(&file_metadata_cache)), + store.id(), + ) .with_coerce_int96(coerce_int96) .with_coerce_int96_tz(coerce_int96_tz.clone()) .fetch_schema_with_location() @@ -428,19 +428,22 @@ impl FileFormat for ParquetFormat { async fn infer_stats( &self, state: &dyn Session, - store: &Arc, + store: &Arc, table_schema: SchemaRef, - object: &ObjectMeta, + object: &FileInfo, ) -> Result { let file_decryption_properties = get_file_decryption_properties(state, &self.options, &object.location) .await?; let file_metadata_cache = state.runtime_env().cache_manager.get_file_metadata_cache(); - DFParquetMetadata::new(store, object) + let reader = store + .open(object, FileAccessContext::new("parquet-statistics")) + .await?; + DFParquetMetadata::new(reader.as_ref(), object) .with_metadata_size_hint(self.metadata_size_hint()) .with_decryption_properties(file_decryption_properties) - .with_file_metadata_cache(Some(file_metadata_cache)) + .with_file_metadata_cache(Some(file_metadata_cache), store.id()) .fetch_statistics(&table_schema) .await } @@ -448,19 +451,22 @@ impl FileFormat for ParquetFormat { async fn infer_ordering( &self, state: &dyn Session, - store: &Arc, + store: &Arc, table_schema: SchemaRef, - object: &ObjectMeta, + object: &FileInfo, ) -> Result> { let file_decryption_properties = get_file_decryption_properties(state, &self.options, &object.location) .await?; let file_metadata_cache = state.runtime_env().cache_manager.get_file_metadata_cache(); - let metadata = DFParquetMetadata::new(store, object) + let reader = store + .open(object, FileAccessContext::new("parquet-statistics")) + .await?; + let metadata = DFParquetMetadata::new(reader.as_ref(), object) .with_metadata_size_hint(self.metadata_size_hint()) .with_decryption_properties(file_decryption_properties) - .with_file_metadata_cache(Some(file_metadata_cache)) + .with_file_metadata_cache(Some(file_metadata_cache), store.id()) .fetch_metadata() .await?; crate::metadata::ordering_from_parquet_metadata(&metadata, &table_schema) @@ -469,19 +475,22 @@ impl FileFormat for ParquetFormat { async fn infer_stats_and_ordering( &self, state: &dyn Session, - store: &Arc, + store: &Arc, table_schema: SchemaRef, - object: &ObjectMeta, + object: &FileInfo, ) -> Result { let file_decryption_properties = get_file_decryption_properties(state, &self.options, &object.location) .await?; let file_metadata_cache = state.runtime_env().cache_manager.get_file_metadata_cache(); - let metadata = DFParquetMetadata::new(store, object) + let reader = store + .open(object, FileAccessContext::new("parquet-statistics")) + .await?; + let metadata = DFParquetMetadata::new(reader.as_ref(), object) .with_metadata_size_hint(self.metadata_size_hint()) .with_decryption_properties(file_decryption_properties) - .with_file_metadata_cache(Some(file_metadata_cache)) + .with_file_metadata_cache(Some(file_metadata_cache), store.id()) .fetch_metadata() .await?; let statistics = DFParquetMetadata::statistics_from_parquet_metadata( @@ -514,15 +523,6 @@ impl FileFormat for ParquetFormat { .ok_or_else(|| internal_datafusion_err!("Expected ParquetSource"))?; source = source.with_table_parquet_options(self.options.clone()); - // Use the CachedParquetFileReaderFactory - let metadata_cache = state.runtime_env().cache_manager.get_file_metadata_cache(); - let store = state - .runtime_env() - .object_store(conf.object_store_url.clone())?; - let cached_parquet_read_factory = - Arc::new(CachedParquetFileReaderFactory::new(store, metadata_cache)); - source = source.with_parquet_file_reader_factory(cached_parquet_read_factory); - if let Some(metadata_size_hint) = metadata_size_hint { source = source.with_metadata_size_hint(metadata_size_hint) } @@ -617,30 +617,6 @@ impl ParquetFormat { } } -/// [`MetadataFetch`] adapter for reading bytes from an [`ObjectStore`] -pub struct ObjectStoreFetch<'a> { - store: &'a dyn ObjectStore, - meta: &'a ObjectMeta, -} - -impl<'a> ObjectStoreFetch<'a> { - pub fn new(store: &'a dyn ObjectStore, meta: &'a ObjectMeta) -> Self { - Self { store, meta } - } -} - -impl MetadataFetch for ObjectStoreFetch<'_> { - fn fetch(&mut self, range: Range) -> BoxFuture<'_, Result> { - async { - self.store - .get_range(&self.meta.location, range) - .await - .map_err(ParquetError::from) - } - .boxed() - } -} - /// Fetches parquet metadata from ObjectStore for given object /// /// This component is a subject to **change** in near future and is exposed for low level integrations @@ -652,17 +628,20 @@ impl MetadataFetch for ObjectStoreFetch<'_> { note = "Use `DFParquetMetadata::fetch_metadata` instead" )] pub async fn fetch_parquet_metadata( - store: &dyn ObjectStore, - object_meta: &ObjectMeta, + store: &StorageBinding, + object_meta: &FileInfo, size_hint: Option, decryption_properties: Option<&FileDecryptionProperties>, file_metadata_cache: Option>, ) -> Result> { let decryption_properties = decryption_properties.cloned().map(Arc::new); - DFParquetMetadata::new(store, object_meta) + let reader = store + .open(object_meta, FileAccessContext::new("parquet-metadata")) + .await?; + DFParquetMetadata::new(reader.as_ref(), object_meta) .with_metadata_size_hint(size_hint) .with_decryption_properties(decryption_properties) - .with_file_metadata_cache(file_metadata_cache) + .with_file_metadata_cache(file_metadata_cache, store.id()) .fetch_metadata() .await } @@ -675,18 +654,21 @@ pub async fn fetch_parquet_metadata( note = "Use `DFParquetMetadata::fetch_statistics` instead" )] pub async fn fetch_statistics( - store: &dyn ObjectStore, + store: &StorageBinding, table_schema: SchemaRef, - file: &ObjectMeta, + file: &FileInfo, metadata_size_hint: Option, decryption_properties: Option<&FileDecryptionProperties>, file_metadata_cache: Option>, ) -> Result { let decryption_properties = decryption_properties.cloned().map(Arc::new); - DFParquetMetadata::new(store, file) + let reader = store + .open(file, FileAccessContext::new("parquet-metadata")) + .await?; + DFParquetMetadata::new(reader.as_ref(), file) .with_metadata_size_hint(metadata_size_hint) .with_decryption_properties(decryption_properties) - .with_file_metadata_cache(file_metadata_cache) + .with_file_metadata_cache(file_metadata_cache, store.id()) .fetch_statistics(&table_schema) .await } diff --git a/datafusion/datasource-parquet/src/metadata.rs b/datafusion/datasource-parquet/src/metadata.rs index 213cac24a85be..0196fd49bf1c5 100644 --- a/datafusion/datasource-parquet/src/metadata.rs +++ b/datafusion/datasource-parquet/src/metadata.rs @@ -15,44 +15,47 @@ // specific language governing permissions and limitations // under the License. -//! [`DFParquetMetadata`] for fetching Parquet file metadata, statistics -//! and schema information. +//! Parquet file metadata, statistics, and schema information. +//! ObjectStore metadata loading is available with the `object_store` feature. -use crate::file_format::ObjectStoreFetch; -use crate::{Int96Coercer, apply_file_schema_type_coercions}; +use crate::Int96Coercer; +use crate::apply_file_schema_type_coercions; +use crate::metadata_io::FileReaderFetch; use arrow::array::{Array, ArrayRef, BooleanArray}; use arrow::compute::kernels::cmp::eq; use arrow::compute::{and, sum}; -use arrow::datatypes::{DataType, Schema, SchemaRef, TimeUnit}; +use arrow::datatypes::TimeUnit; +use arrow::datatypes::{DataType, Schema, SchemaRef}; use datafusion_common::encryption::FileDecryptionProperties; use datafusion_common::stats::Precision; -use datafusion_common::{ - ColumnStatistics, DataFusionError, HashMap, Result, ScalarValue, Statistics, - internal_datafusion_err, -}; -use datafusion_execution::cache::cache_manager::{ - CachedFileMetadataEntry, FileMetadata, FileMetadataCache, -}; +use datafusion_common::{ColumnStatistics, HashMap, Result, ScalarValue, Statistics}; +use datafusion_execution::cache::cache_manager::FileMetadata; use datafusion_functions_aggregate_common::min_max::{MaxAccumulator, MinAccumulator}; use datafusion_physical_expr::expressions::Column; use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; use datafusion_physical_plan::Accumulator; +use datafusion_storage::FileInfo; +use datafusion_storage::FileReader; +use datafusion_storage::path::Path; use log::debug; -use object_store::path::Path; -use object_store::{ObjectMeta, ObjectStore}; -use parquet::DecodeResult; use parquet::arrow::arrow_reader::statistics::StatisticsConverter; use parquet::arrow::{parquet_column, parquet_to_arrow_schema}; use parquet::basic::{ColumnOrder, SortOrder, Type as PhysicalType}; -use parquet::file::metadata::{ - PageIndexPolicy, ParquetMetaData, ParquetMetaDataPushDecoder, ParquetMetaDataReader, - RowGroupMetaData, SortingColumn, -}; +use parquet::file::metadata::{ParquetMetaData, RowGroupMetaData, SortingColumn}; use parquet::file::statistics::Statistics as ParquetStatistics; use parquet::schema::types::{ColumnDescriptor, SchemaDescriptor}; use std::any::Any; use std::sync::Arc; +use datafusion_execution::cache::cache_manager::{ + CachedFileMetadataEntry, FileMetadataCache, +}; +use parquet::file::metadata::{ + PageIndexPolicy, ParquetMetaDataPushDecoder, ParquetMetaDataReader, +}; + +use datafusion_common::{DataFusionError, internal_datafusion_err}; + /// Minimum fraction of row groups that must report NDV statistics for the /// merged result to be `Inexact` rather than `Absent`, as the estimate /// would be too unreliable otherwise. @@ -121,10 +124,11 @@ pub(crate) fn has_untrusted_byte_array_stats<'a>( /// [`ParquetFileReaderFactory`]: crate::ParquetFileReaderFactory #[derive(Debug)] pub struct DFParquetMetadata<'a> { + cache_key: Option, /// Source of the Parquet file's bytes. - store: &'a dyn ObjectStore, + store: &'a dyn FileReader, /// Location, size and last-modified time of the target Parquet file. - object_meta: &'a ObjectMeta, + object_meta: &'a FileInfo, /// Hint for the number of trailing bytes to prefetch before parsing the /// footer, mirroring [`ParquetMetaDataReader::with_prefetch_hint`]. metadata_size_hint: Option, @@ -153,8 +157,9 @@ impl<'a> DFParquetMetadata<'a> { /// /// Use the `with_*` builder methods to customize behavior /// before calling [`Self::fetch_metadata`] or [`Self::fetch_schema`]. - pub fn new(store: &'a dyn ObjectStore, object_meta: &'a ObjectMeta) -> Self { + pub fn new(store: &'a dyn FileReader, object_meta: &'a FileInfo) -> Self { Self { + cache_key: None, store, object_meta, metadata_size_hint: None, @@ -196,7 +201,12 @@ impl<'a> DFParquetMetadata<'a> { pub fn with_file_metadata_cache( mut self, file_metadata_cache: Option>, + storage_id: u64, ) -> Self { + self.cache_key = Some(datafusion_execution::cache::FileCacheKey { + storage_id, + path: self.object_meta.location.clone(), + }); self.file_metadata_cache = file_metadata_cache; self } @@ -265,7 +275,8 @@ impl<'a> DFParquetMetadata<'a> { if cache_metadata && let Some(file_metadata_cache) = self.file_metadata_cache.as_ref() - && let Some(cached) = file_metadata_cache.get(&self.object_meta.location) + && let Some(cached) = + file_metadata_cache.get(self.cache_key.as_ref().expect("cache scope")) && cached.is_valid_for(self.object_meta) && let Some(cached_parquet) = cached .file_metadata @@ -329,7 +340,7 @@ impl<'a> DFParquetMetadata<'a> { fn cache_metadata(&self, metadata: Arc) -> Result<()> { if let Some(file_metadata_cache) = &self.file_metadata_cache { file_metadata_cache.put( - &self.object_meta.location, + self.cache_key.as_ref().expect("cache scope"), CachedFileMetadataEntry::new( self.object_meta.clone(), Arc::new(CachedParquetMetaData::new(metadata)), @@ -357,52 +368,25 @@ impl<'a> DFParquetMetadata<'a> { decoder = decoder.with_page_index_policy(page_index_policy); - if let Some(hint) = self.metadata_size_hint { - let prefetch_start = file_size.saturating_sub(hint as u64); - let prefetch_range = prefetch_start..file_size; - let data = self - .store - .get_ranges( - &self.object_meta.location, - std::slice::from_ref(&prefetch_range), - ) - .await - .map_err(DataFusionError::from)?; - decoder - .push_ranges(vec![prefetch_range], data) - .map_err(DataFusionError::from)?; - } - - let metadata = loop { - match decoder.try_decode().map_err(DataFusionError::from)? { - DecodeResult::Data(metadata) => break metadata, - DecodeResult::NeedsData(ranges) => { - let buffers = self - .store - .get_ranges(&self.object_meta.location, &ranges) - .await - .map_err(DataFusionError::from)?; - decoder - .push_ranges(ranges, buffers) - .map_err(DataFusionError::from)?; - } - DecodeResult::Finished => { - return Err(DataFusionError::Internal( - "ParquetMetaDataPushDecoder finished without producing metadata" - .to_string(), - )); - } - } - }; - - Ok(Arc::new(metadata)) + crate::metadata_io::decode_metadata( + decoder, + file_size, + self.metadata_size_hint, + |ranges| async move { + self.store + .read_ranges(ranges) + .await + .map_err(DataFusionError::from) + }, + ) + .await } /// If `metadata` does not already have a page index, fetch and attach the /// column and offset indexes. async fn load_page_index( - store: &dyn ObjectStore, - object_meta: &ObjectMeta, + store: &dyn FileReader, + _object_meta: &FileInfo, metadata: Arc, ) -> Result> { if metadata.column_index().is_some() && metadata.offset_index().is_some() { @@ -412,7 +396,7 @@ impl<'a> DFParquetMetadata<'a> { Arc::try_unwrap(metadata).unwrap_or_else(|shared| (*shared).clone()); let mut reader = ParquetMetaDataReader::new_with_metadata(metadata) .with_page_index_policy(PageIndexPolicy::Optional); - let fetch = ObjectStoreFetch::new(store, object_meta); + let fetch = FileReaderFetch::new(store); reader .load_page_index(fetch) .await @@ -496,152 +480,158 @@ impl<'a> DFParquetMetadata<'a> { /// - If neither method is applicable, byte size is marked as Precision::Absent pub fn statistics_from_parquet_metadata( metadata: &ParquetMetaData, - logical_file_schema: &SchemaRef, + table_schema: &SchemaRef, ) -> Result { - let row_groups_metadata = metadata.row_groups(); - - // Use Statistics::default() as opposed to Statistics::new_unknown() - // because we are going to replace the column statistics below - // and we don't want to initialize them twice. - let mut statistics = Statistics::default(); - let mut has_statistics = false; - let mut num_rows = 0_usize; - for row_group_meta in row_groups_metadata { - num_rows += row_group_meta.num_rows() as usize; - - if !has_statistics { - has_statistics = row_group_meta - .columns() - .iter() - .any(|column| column.statistics().is_some()); - } - } - statistics.num_rows = Precision::Exact(num_rows); - - let file_metadata = metadata.file_metadata(); - let mut physical_file_schema = parquet_to_arrow_schema( - file_metadata.schema_descr(), - file_metadata.key_value_metadata(), - )?; + statistics_from_parquet_metadata(metadata, table_schema) + } +} - if let Some(merged) = - apply_file_schema_type_coercions(logical_file_schema, &physical_file_schema) - { - physical_file_schema = merged; +pub fn statistics_from_parquet_metadata( + metadata: &ParquetMetaData, + logical_file_schema: &SchemaRef, +) -> Result { + let row_groups_metadata = metadata.row_groups(); + + // Use Statistics::default() as opposed to Statistics::new_unknown() + // because we are going to replace the column statistics below + // and we don't want to initialize them twice. + let mut statistics = Statistics::default(); + let mut has_statistics = false; + let mut num_rows = 0_usize; + for row_group_meta in row_groups_metadata { + num_rows += row_group_meta.num_rows() as usize; + + if !has_statistics { + has_statistics = row_group_meta + .columns() + .iter() + .any(|column| column.statistics().is_some()); } + } + statistics.num_rows = Precision::Exact(num_rows); + + let file_metadata = metadata.file_metadata(); + let mut physical_file_schema = parquet_to_arrow_schema( + file_metadata.schema_descr(), + file_metadata.key_value_metadata(), + )?; + + if let Some(merged) = + apply_file_schema_type_coercions(logical_file_schema, &physical_file_schema) + { + physical_file_schema = merged; + } - statistics.column_statistics = - if has_statistics { - let (mut max_accs, mut min_accs) = - create_max_min_accs(logical_file_schema); - let mut null_counts_array = - vec![Precision::Absent; logical_file_schema.fields().len()]; - let mut column_byte_sizes = - vec![Precision::Absent; logical_file_schema.fields().len()]; - let mut is_max_value_exact = - vec![Some(true); logical_file_schema.fields().len()]; - let mut is_min_value_exact = - vec![Some(true); logical_file_schema.fields().len()]; - let mut distinct_counts_array = - vec![Precision::Absent; logical_file_schema.fields().len()]; - logical_file_schema.fields().iter().enumerate().for_each( - |(idx, field)| match StatisticsConverter::try_new( - field.name(), - &physical_file_schema, - file_metadata.schema_descr(), - ) { - Ok(stats_converter) => { - let parquet_index = stats_converter.parquet_column_index(); - if parquet_index.is_some_and(|index| { - has_untrusted_min_max_order( - file_metadata.schema_descr(), - file_metadata.column_orders().map(Vec::as_slice), - index, - ) - }) || has_untrusted_byte_array_stats( + statistics.column_statistics = if has_statistics { + let (mut max_accs, mut min_accs) = create_max_min_accs(logical_file_schema); + let mut null_counts_array = + vec![Precision::Absent; logical_file_schema.fields().len()]; + let mut column_byte_sizes = + vec![Precision::Absent; logical_file_schema.fields().len()]; + let mut is_max_value_exact = vec![Some(true); logical_file_schema.fields().len()]; + let mut is_min_value_exact = vec![Some(true); logical_file_schema.fields().len()]; + let mut distinct_counts_array = + vec![Precision::Absent; logical_file_schema.fields().len()]; + logical_file_schema + .fields() + .iter() + .enumerate() + .for_each(|(idx, field)| { + match StatisticsConverter::try_new( + field.name(), + &physical_file_schema, + file_metadata.schema_descr(), + ) { + Ok(stats_converter) => { + let parquet_index = stats_converter.parquet_column_index(); + if parquet_index.is_some_and(|index| { + has_untrusted_min_max_order( file_metadata.schema_descr(), - parquet_index, - row_groups_metadata, - ) { - // The remaining row groups cannot establish bounds - // for the whole file. Keep unrelated statistics. - min_accs[idx] = None; - max_accs[idx] = None; - } - let mut accumulators = StatisticsAccumulators { - min_accs: &mut min_accs, - max_accs: &mut max_accs, - null_counts_array: &mut null_counts_array, - is_min_value_exact: &mut is_min_value_exact, - is_max_value_exact: &mut is_max_value_exact, - column_byte_sizes: &mut column_byte_sizes, - distinct_counts_array: &mut distinct_counts_array, - }; - summarize_column_statistics( - logical_file_schema, - &mut accumulators, - idx, - &stats_converter, - row_groups_metadata, - num_rows, + file_metadata.column_orders().map(Vec::as_slice), + index, ) - .ok(); - } - Err(e) => { - debug!("Failed to create statistics converter: {e}"); - null_counts_array[idx] = Precision::Exact(num_rows); - } - }, - ); - - let mut accumulators = StatisticsAccumulators { - min_accs: &mut min_accs, - max_accs: &mut max_accs, - null_counts_array: &mut null_counts_array, - is_min_value_exact: &mut is_min_value_exact, - is_max_value_exact: &mut is_max_value_exact, - column_byte_sizes: &mut column_byte_sizes, - distinct_counts_array: &mut distinct_counts_array, - }; - accumulators.build_column_statistics(logical_file_schema) - } else { - // Record column sizes - logical_file_schema - .fields() - .iter() - .enumerate() - .map(|(logical_file_schema_index, field)| { - let arrow_field = - logical_file_schema.field(logical_file_schema_index); - let parquet_idx = parquet_column( + }) || has_untrusted_byte_array_stats( file_metadata.schema_descr(), - &physical_file_schema, - arrow_field.name(), - ) - .map(|(idx, _)| idx); - let byte_size = compute_arrow_column_size( - field.data_type(), + parquet_index, + row_groups_metadata, + ) { + // The remaining row groups cannot establish bounds + // for the whole file. Keep unrelated statistics. + min_accs[idx] = None; + max_accs[idx] = None; + } + let mut accumulators = StatisticsAccumulators { + min_accs: &mut min_accs, + max_accs: &mut max_accs, + null_counts_array: &mut null_counts_array, + is_min_value_exact: &mut is_min_value_exact, + is_max_value_exact: &mut is_max_value_exact, + column_byte_sizes: &mut column_byte_sizes, + distinct_counts_array: &mut distinct_counts_array, + }; + summarize_column_statistics( + logical_file_schema, + &mut accumulators, + idx, + &stats_converter, row_groups_metadata, - parquet_idx, num_rows, - ); - ColumnStatistics::new_unknown().with_byte_size(byte_size) - }) - .collect() - }; + ) + .ok(); + } + Err(e) => { + debug!("Failed to create statistics converter: {e}"); + null_counts_array[idx] = Precision::Exact(num_rows); + } + } + }); - #[cfg(debug_assertions)] - { - // Check that the column statistics length matches the table schema fields length - assert_eq!( - statistics.column_statistics.len(), - logical_file_schema.fields().len(), - "Column statistics length does not match table schema fields length" - ); - } + let mut accumulators = StatisticsAccumulators { + min_accs: &mut min_accs, + max_accs: &mut max_accs, + null_counts_array: &mut null_counts_array, + is_min_value_exact: &mut is_min_value_exact, + is_max_value_exact: &mut is_max_value_exact, + column_byte_sizes: &mut column_byte_sizes, + distinct_counts_array: &mut distinct_counts_array, + }; + accumulators.build_column_statistics(logical_file_schema) + } else { + // Record column sizes + logical_file_schema + .fields() + .iter() + .enumerate() + .map(|(logical_file_schema_index, field)| { + let arrow_field = logical_file_schema.field(logical_file_schema_index); + let parquet_idx = parquet_column( + file_metadata.schema_descr(), + &physical_file_schema, + arrow_field.name(), + ) + .map(|(idx, _)| idx); + let byte_size = compute_arrow_column_size( + field.data_type(), + row_groups_metadata, + parquet_idx, + num_rows, + ); + ColumnStatistics::new_unknown().with_byte_size(byte_size) + }) + .collect() + }; - Ok(statistics) + #[cfg(debug_assertions)] + { + // Check that the column statistics length matches the table schema fields length + assert_eq!( + statistics.column_statistics.len(), + logical_file_schema.fields().len(), + "Column statistics length does not match table schema fields length" + ); } + + Ok(statistics) } /// Min/max aggregation can take Dictionary encode input but always produces unpacked diff --git a/datafusion/datasource-parquet/src/metadata_io.rs b/datafusion/datasource-parquet/src/metadata_io.rs new file mode 100644 index 0000000000000..d6d498f0ee46c --- /dev/null +++ b/datafusion/datasource-parquet/src/metadata_io.rs @@ -0,0 +1,76 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Shared metadata decoder driver for all file access implementations. +use bytes::Bytes; +use datafusion_common::{DataFusionError, Result}; +use parquet::DecodeResult; +use parquet::file::metadata::{ParquetMetaData, ParquetMetaDataPushDecoder}; +use std::{future::Future, ops::Range, sync::Arc}; + +pub(crate) async fn decode_metadata( + mut decoder: ParquetMetaDataPushDecoder, + file_size: u64, + hint: Option, + mut read: F, +) -> Result> +where + F: FnMut(Vec>) -> Fut, + Fut: Future>>, +{ + if let Some(hint) = hint.filter(|hint| *hint > 0) { + let range = file_size.saturating_sub(hint as u64)..file_size; + let ranges = vec![range]; + let buffers = read(ranges.clone()).await?; + decoder.push_ranges(ranges, buffers)?; + } + loop { + match decoder.try_decode()? { + DecodeResult::Data(metadata) => return Ok(Arc::new(metadata)), + DecodeResult::NeedsData(ranges) => { + let buffers = read(ranges.clone()).await?; + decoder.push_ranges(ranges, buffers)?; + } + DecodeResult::Finished => { + return Err(DataFusionError::Internal( + "Parquet metadata decoder finished without producing metadata".into(), + )); + } + } + } +} + +/// Adapt a file reader to the Parquet page-index fetch interface. +pub(crate) struct FileReaderFetch<'a>(&'a dyn datafusion_storage::FileReader); +impl<'a> FileReaderFetch<'a> { + pub(crate) fn new(reader: &'a dyn datafusion_storage::FileReader) -> Self { + Self(reader) + } +} +impl parquet::arrow::async_reader::MetadataFetch for FileReaderFetch<'_> { + fn fetch( + &mut self, + range: Range, + ) -> futures::future::BoxFuture<'_, parquet::errors::Result> { + Box::pin(async move { + self.0 + .read_range(range.into()) + .await + .map_err(|e| parquet::errors::ParquetError::External(Box::new(e))) + }) + } +} diff --git a/datafusion/datasource-parquet/src/mod.rs b/datafusion/datasource-parquet/src/mod.rs index 23348f94aae8f..203c16828539f 100644 --- a/datafusion/datasource-parquet/src/mod.rs +++ b/datafusion/datasource-parquet/src/mod.rs @@ -29,6 +29,7 @@ mod bloom_filter; mod decoder_projection; pub mod file_format; pub mod metadata; +mod metadata_io; mod metrics; mod nested_schema_pruning; mod opener; @@ -44,6 +45,7 @@ mod sort; pub mod source; #[cfg(test)] mod statistics_order_tests; +pub mod storage; mod supported_predicates; #[cfg(test)] mod test_util; diff --git a/datafusion/datasource-parquet/src/opener/encryption.rs b/datafusion/datasource-parquet/src/opener/encryption.rs index 498fe8acf7530..4cbc4d06ebb10 100644 --- a/datafusion/datasource-parquet/src/opener/encryption.rs +++ b/datafusion/datasource-parquet/src/opener/encryption.rs @@ -58,7 +58,7 @@ impl EncryptionContext { pub(super) async fn get_file_decryption_properties( &self, - file_location: &object_store::path::Path, + file_location: &datafusion_storage::path::Path, ) -> Result>> { match &self.file_decryption_properties { Some(file_decryption_properties) => { @@ -80,7 +80,7 @@ impl EncryptionContext { impl EncryptionContext { pub(super) async fn get_file_decryption_properties( &self, - _file_location: &object_store::path::Path, + _file_location: &datafusion_storage::path::Path, ) -> Result< Option>, > { diff --git a/datafusion/datasource-parquet/src/opener/mod.rs b/datafusion/datasource-parquet/src/opener/mod.rs index 25c3bc9a77851..0dabf21d74303 100644 --- a/datafusion/datasource-parquet/src/opener/mod.rs +++ b/datafusion/datasource-parquet/src/opener/mod.rs @@ -2306,7 +2306,9 @@ mod test { parquet_file_reader_factory: self .parquet_file_reader_factory .unwrap_or_else(|| { - Arc::new(DefaultParquetFileReaderFactory::new(store)) as _ + Arc::new(DefaultParquetFileReaderFactory::new( + test_utils::storage::object_store(store), + )) as _ }), pushdown_filters: self.pushdown_filters, reorder_filters: self.reorder_filters, @@ -3671,10 +3673,12 @@ mod test { use parquet::file::properties::WriterProperties; let store = Arc::new(InMemory::new()) as Arc; + let binding = test_utils::storage::object_store(store.clone()); let metadata_cache: Arc = - Arc::new(DefaultCache::::new( - 64 * 1024 * 1024, - )); + Arc::new(DefaultCache::< + datafusion_execution::cache::FileCacheKey, + CachedFileMetadataEntry, + >::new(64 * 1024 * 1024)); let values: Vec = (1..=100).collect(); let batch = record_batch!(( "a", @@ -3711,7 +3715,7 @@ mod test { .with_metrics(metrics.clone()) .with_parquet_file_reader_factory(Arc::new( CachedParquetFileReaderFactory::new( - Arc::clone(&store), + Arc::clone(&binding), Arc::clone(&metadata_cache), ), )) @@ -3723,7 +3727,10 @@ mod test { assert_eq!(counter_metric_value(&metrics, "page_index_load_skipped"), 1); let cached = metadata_cache - .get(&Path::from("test.parquet")) + .get(&datafusion_execution::cache::FileCacheKey { + storage_id: binding.id(), + path: datafusion_storage::path::Path::from("test.parquet"), + }) .expect("metadata cache should contain the file"); let extra_info = cached.file_metadata.extra_info(); let page_index_cached = extra_info.get("page_index").map(String::as_str); diff --git a/datafusion/datasource-parquet/src/reader.rs b/datafusion/datasource-parquet/src/reader.rs index 71b0020f32f64..f71d40ac531a0 100644 --- a/datafusion/datasource-parquet/src/reader.rs +++ b/datafusion/datasource-parquet/src/reader.rs @@ -15,8 +15,7 @@ // specific language governing permissions and limitations // under the License. -//! [`ParquetFileReaderFactory`] and [`DefaultParquetFileReaderFactory`] for -//! low level control of parquet file readers +//! [`ParquetFileReaderFactory`] for low level control of Parquet file readers. use crate::ParquetFileMetrics; use crate::metadata::DFParquetMetadata; @@ -26,10 +25,9 @@ use datafusion_datasource::PartitionedFile; use datafusion_execution::cache::cache_manager::FileMetadata; use datafusion_execution::cache::cache_manager::FileMetadataCache; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; +use datafusion_storage::{FileAccessContext, FileReader, StorageBinding}; use futures::FutureExt; -use futures::TryFutureExt; use futures::future::BoxFuture; -use object_store::{ObjectStore, ObjectStoreExt}; use parquet::arrow::arrow_reader::ArrowReaderOptions; use parquet::arrow::async_reader::AsyncFileReader; use parquet::errors::ParquetError; @@ -38,6 +36,7 @@ use std::any::Any; use std::fmt::Debug; use std::ops::Range; use std::sync::Arc; +use tokio::sync::OnceCell; /// Interface for reading Apache Parquet files. /// @@ -45,7 +44,7 @@ use std::sync::Arc; /// [`AsyncFileReader`] can be used to provide custom data access operations /// such as pre-cached metadata, I/O coalescing, etc. /// -/// See [`DefaultParquetFileReaderFactory`] for a simple implementation. +/// See [`DefaultParquetFileReaderFactory`] for the built-in implementation. pub trait ParquetFileReaderFactory: Debug + Send + Sync + 'static { /// Provides an `AsyncFileReader` for reading data from a parquet file specified /// @@ -73,18 +72,27 @@ pub trait ParquetFileReaderFactory: Debug + Send + Sync + 'static { /// Default implementation of [`ParquetFileReaderFactory`] /// /// This implementation: -/// 1. Reads parquet directly from an underlying [`ObjectStore`] instance. +/// 1. Reads parquet directly from an underlying [`StorageBinding`] instance. /// 2. Reads the footer and page metadata on demand. /// 3. Does not cache metadata or coalesce I/O operations. #[derive(Debug)] pub struct DefaultParquetFileReaderFactory { - store: Arc, + store: Arc, + context: FileAccessContext, } impl DefaultParquetFileReaderFactory { + pub fn with_context(mut self, context: FileAccessContext) -> Self { + self.context = context; + self + } + /// Create a new `DefaultParquetFileReaderFactory`. - pub fn new(store: Arc) -> Self { - Self { store } + pub fn new(store: Arc) -> Self { + Self { + store, + context: FileAccessContext::new("parquet"), + } } } @@ -107,6 +115,7 @@ impl ParquetFileReaderFactory for DefaultParquetFileReaderFactory { Arc::clone(&self.store), partitioned_file, ) + .with_context(self.context.clone()) .with_metadata_hint(metadata_size_hint); Ok(Box::new(reader)) } @@ -120,17 +129,24 @@ impl ParquetFileReaderFactory for DefaultParquetFileReaderFactory { /// parquet opener can skip page-index I/O during the initial metadata load. #[derive(Debug)] pub struct CachedParquetFileReaderFactory { - store: Arc, + store: Arc, + context: FileAccessContext, metadata_cache: Arc, } impl CachedParquetFileReaderFactory { + pub fn with_context(mut self, context: FileAccessContext) -> Self { + self.context = context; + self + } + pub fn new( - store: Arc, + store: Arc, metadata_cache: Arc, ) -> Self { Self { store, + context: FileAccessContext::new("parquet"), metadata_cache, } } @@ -155,6 +171,7 @@ impl ParquetFileReaderFactory for CachedParquetFileReaderFactory { Arc::clone(&self.store), partitioned_file, ) + .with_context(self.context.clone()) .with_metadata_hint(metadata_size_hint) .with_metadata_cache(Some(Arc::clone(&self.metadata_cache))); @@ -164,7 +181,7 @@ impl ParquetFileReaderFactory for CachedParquetFileReaderFactory { /// Implements [`AsyncFileReader`] for a parquet file in object storage. /// -/// This implementation reads data directly from the underlying [`ObjectStore`] +/// This implementation reads data directly from the underlying [`StorageBinding`] /// on demand, as required, tracking the number of bytes read. /// /// When configured via [`Self::with_metadata_cache`], [`Self::get_metadata`] @@ -174,17 +191,33 @@ impl ParquetFileReaderFactory for CachedParquetFileReaderFactory { /// # Notes /// /// This implementation does not coalesce I/O operations or cache bytes. Such -/// optimizations can be done either at the object store level or by providing +/// optimizations can be done either at the storage backend level or by providing /// a custom implementation of [`ParquetFileReaderFactory`]. pub struct ParquetFileReader { + reader: OnceCell>, file_metrics: ParquetFileMetrics, - store: Arc, + store: Arc, + context: FileAccessContext, partitioned_file: PartitionedFile, metadata_cache: Option>, metadata_size_hint: Option, } impl ParquetFileReader { + fn with_context(mut self, context: FileAccessContext) -> Self { + self.context = context; + self + } + async fn reader(&self) -> parquet::errors::Result<&Arc> { + self.reader + .get_or_try_init(|| { + self.store + .open(&self.partitioned_file.object_meta, self.context.clone()) + }) + .await + .map_err(|e| ParquetError::External(Box::new(e))) + } + /// Create a new `ParquetFileReader`. /// /// By default the reader has no [`FileMetadataCache`] and no metadata @@ -195,12 +228,14 @@ impl ParquetFileReader { /// [`Self::with_metadata_hint`] to set the size hint. pub(crate) fn new( file_metrics: ParquetFileMetrics, - store: Arc, + store: Arc, partitioned_file: PartitionedFile, ) -> Self { Self { + reader: OnceCell::new(), file_metrics, store, + context: FileAccessContext::new("parquet"), partitioned_file, metadata_cache: None, metadata_size_hint: None, @@ -242,10 +277,14 @@ impl AsyncFileReader for ParquetFileReader { ) -> BoxFuture<'_, parquet::errors::Result> { let bytes_scanned = range.end - range.start; self.file_metrics.bytes_scanned.add(bytes_scanned as usize); - self.store - .get_range(&self.partitioned_file.object_meta.location, range) - .map_err(|e| ParquetError::External(Box::new(e))) - .boxed() + async move { + self.reader() + .await? + .read_range(range.into()) + .await + .map_err(|e| ParquetError::External(Box::new(e))) + } + .boxed() } fn get_byte_ranges( @@ -258,8 +297,9 @@ impl AsyncFileReader for ParquetFileReader { let total: u64 = ranges.iter().map(|r| r.end - r.start).sum(); self.file_metrics.bytes_scanned.add(total as usize); async move { - self.store - .get_ranges(&self.partitioned_file.object_meta.location, &ranges) + self.reader() + .await? + .read_ranges(ranges) .await .map_err(|e| ParquetError::External(Box::new(e))) } @@ -284,9 +324,9 @@ impl AsyncFileReader for ParquetFileReader { let page_index_policy = options.map(|o| o.column_index_policy()); - DFParquetMetadata::new(&self.store, &object_meta) + DFParquetMetadata::new(self.reader().await?.as_ref(), &object_meta) .with_decryption_properties(file_decryption_properties) - .with_file_metadata_cache(metadata_cache) + .with_file_metadata_cache(metadata_cache, self.store.id()) .with_metadata_size_hint(self.metadata_size_hint) .with_page_index_policy(page_index_policy) .fetch_metadata() diff --git a/datafusion/datasource-parquet/src/schema_coercion.rs b/datafusion/datasource-parquet/src/schema_coercion.rs index 825b7e27ff20d..1107b67362ecf 100644 --- a/datafusion/datasource-parquet/src/schema_coercion.rs +++ b/datafusion/datasource-parquet/src/schema_coercion.rs @@ -429,7 +429,7 @@ fn field_with_new_type(field: &FieldRef, new_type: DataType) -> FieldRef { /// Transform a schema to use view types for Utf8 and Binary /// -/// See [`ParquetFormat::force_view_types`](crate::file_format::ParquetFormat::force_view_types) for details +/// See [`schema_force_view_types`](datafusion_common::config::ParquetOptions::schema_force_view_types) for details. pub fn transform_schema_to_view(schema: &Schema) -> Schema { let transformed_fields: Vec> = schema .fields diff --git a/datafusion/datasource-parquet/src/sink.rs b/datafusion/datasource-parquet/src/sink.rs index 3c66d4dcd74fb..3fa68d51e3831 100644 --- a/datafusion/datasource-parquet/src/sink.rs +++ b/datafusion/datasource-parquet/src/sink.rs @@ -16,7 +16,7 @@ // under the License. //! [`ParquetSink`] — DataFusion `DataSink` implementation that writes one -//! or more Parquet files to an [`ObjectStore`], optionally with parallel +//! or more Parquet files to a [`StorageBinding`], optionally with parallel //! per-column and per-row-group serialization. use std::fmt; @@ -36,9 +36,7 @@ use datafusion_datasource::sink::DataSink; #[cfg(feature = "proto")] use datafusion_datasource::sink::DataSinkExec; use datafusion_datasource::write::demux::DemuxedStreamReceiver; -use datafusion_datasource::write::{ - ObjectWriterBuilder, SharedBuffer, get_writer_schema, -}; +use datafusion_datasource::write::{FileWriterBuilder, SharedBuffer, get_writer_schema}; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryPool, MemoryReservation}; use datafusion_execution::runtime_env::RuntimeEnv; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; @@ -49,9 +47,9 @@ use datafusion_physical_plan::metrics::{ MetricsSet, Time, }; use datafusion_physical_plan::{DisplayAs, DisplayFormatType}; -use object_store::ObjectStore; -use object_store::buffered::BufWriter; -use object_store::path::Path; +use datafusion_storage::StorageBinding; +use datafusion_storage::path::Path; +use datafusion_storage::{FileAccessContext, FileOutput, WriterOptions}; use parquet::arrow::arrow_writer::{ ArrowColumnChunk, ArrowColumnWriter, ArrowLeafColumn, ArrowRowGroupWriterFactory, ArrowWriterOptions, compute_leaves, @@ -64,6 +62,7 @@ use parquet::file::properties::{ DEFAULT_MAX_ROW_GROUP_ROW_COUNT, WriterProperties, WriterPropertiesBuilder, }; use parquet::file::writer::SerializedFileWriter; + use tokio::io::{AsyncWrite, AsyncWriteExt}; use tokio::sync::mpsc::{self, Receiver, Sender}; @@ -175,22 +174,28 @@ impl ParquetSink { /// Creates an AsyncArrowWriter which serializes a parquet file to an ObjectStore /// AsyncArrowWriters are used when individual parquet file serialization is not parallelized - fn create_async_arrow_writer( + async fn create_async_arrow_writer( &self, location: &Path, - object_store: Arc, + object_store: Arc, context: &Arc, parquet_props: WriterProperties, - ) -> Result> { - let buf_writer = BufWriter::with_capacity( - object_store, - location.clone(), - context - .session_config() - .options() - .execution - .objectstore_writer_buffer_size, - ); + ) -> Result> { + let buf_writer = object_store + .writer( + location, + WriterOptions { + buffer_size: Some( + context + .session_config() + .options() + .execution + .objectstore_writer_buffer_size, + ), + }, + FileAccessContext::new("parquet-write"), + ) + .await?; let options = ArrowWriterOptions::new() .with_properties(parquet_props) .with_skip_arrow_metadata(self.parquet_options.global.skip_arrow_metadata); @@ -265,7 +270,7 @@ impl FileSink for ParquetSink { context: &Arc, demux_task: SpawnedTask>, mut file_stream_rx: DemuxedStreamReceiver, - object_store: Arc, + object_store: Arc, ) -> Result { let rows_written_counter = MetricBuilder::new(&self.metrics) .with_category(MetricCategory::Rows) @@ -300,12 +305,14 @@ impl FileSink for ParquetSink { if !parquet_opts.global.allow_single_file_parallelism || parquet_opts.global.content_defined_chunking.enabled { - let mut writer = self.create_async_arrow_writer( - &path, - Arc::clone(&object_store), - context, - parquet_props.clone(), - )?; + let mut writer = self + .create_async_arrow_writer( + &path, + Arc::clone(&object_store), + context, + parquet_props.clone(), + ) + .await?; let reservation = MemoryConsumer::new(format!("ParquetSink[{path}]")) .register(context.memory_pool()); file_write_tasks.spawn( @@ -323,7 +330,7 @@ impl FileSink for ParquetSink { .with_elapsed_compute(elapsed_compute.clone()), ); } else { - let writer = ObjectWriterBuilder::new( + let writer = FileWriterBuilder::new( // Parquet files as a whole are never compressed, since they // manage compressed blocks themselves. FileCompressionType::UNCOMPRESSED, @@ -337,7 +344,8 @@ impl FileSink for ParquetSink { .execution .objectstore_writer_buffer_size, )) - .build()?; + .build() + .await?; let ctx = ParquetFileWriteContext { schema: get_writer_schema(&self.config), props: Arc::new(parquet_props), @@ -499,7 +507,12 @@ impl ParquetSink { "ParquetSinkExecNode is missing required field 'sink'" ) })?; - let data_sink = ParquetSink::try_from(proto_sink)?; + let mut data_sink = ParquetSink::try_from(proto_sink)?; + data_sink.config.storage = Some( + ctx.task_ctx() + .runtime_env() + .storage(&data_sink.config.object_store_url)?, + ); let sort_order = DataSinkExec::decode_sort_order( sink_node.sort_order.as_ref(), ctx, @@ -759,7 +772,7 @@ fn spawn_parquet_parallel_serialization_task( } /// Consume RowGroups serialized by other parallel tasks and concatenate them in -/// to the final parquet file, while flushing finalized bytes to an [ObjectStore] +/// to the final parquet file, while flushing finalized bytes to a [`StorageBinding`] async fn concatenate_parallel_row_groups( mut parquet_writer: SerializedFileWriter, merged_buff: SharedBuffer, diff --git a/datafusion/datasource-parquet/src/source.rs b/datafusion/datasource-parquet/src/source.rs index 4872db9fd3329..2e42d5c2f91ae 100644 --- a/datafusion/datasource-parquet/src/source.rs +++ b/datafusion/datasource-parquet/src/source.rs @@ -20,7 +20,6 @@ use std::fmt::Debug; use std::fmt::Formatter; use std::sync::Arc; -use crate::DefaultParquetFileReaderFactory; use crate::ParquetFileReaderFactory; use crate::opener::ParquetMorselizer; use crate::opener::build_pruning_predicates; @@ -71,8 +70,8 @@ use log::warn; #[cfg(feature = "parquet_encryption")] use datafusion_execution::parquet_encryption::EncryptionFactory; use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; +use datafusion_storage::StorageBinding; use itertools::Itertools; -use object_store::ObjectStore; use parquet::arrow::RowNumber; #[cfg(feature = "parquet_encryption")] use parquet::encryption::decrypt::FileDecryptionProperties; @@ -112,13 +111,13 @@ use parquet::encryption::decrypt::FileDecryptionProperties; /// # use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; /// # use datafusion_datasource_parquet::source::ParquetSource; /// # use datafusion_datasource::PartitionedFile; -/// # use datafusion_execution::object_store::ObjectStoreUrl; +/// # use datafusion_storage::StorageUrl; /// # use datafusion_physical_expr::expressions::lit; /// # use datafusion_datasource::source::DataSourceExec; /// # use datafusion_common::config::TableParquetOptions; /// /// # let file_schema = Arc::new(Schema::empty()); -/// # let object_store_url = ObjectStoreUrl::local_filesystem(); +/// # let object_store_url = StorageUrl::local_filesystem(); /// # let predicate = lit(true); /// let source = Arc::new( /// ParquetSource::new(Arc::clone(&file_schema)) @@ -240,7 +239,7 @@ use parquet::encryption::decrypt::FileDecryptionProperties; /// # use datafusion_datasource_parquet::ParquetAccessPlan; /// # use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; /// # use datafusion_datasource_parquet::source::ParquetSource; -/// # use datafusion_execution::object_store::ObjectStoreUrl; +/// # use datafusion_storage::StorageUrl; /// # use datafusion_datasource::source::DataSourceExec; /// /// # fn schema() -> SchemaRef { @@ -254,7 +253,7 @@ use parquet::encryption::decrypt::FileDecryptionProperties; /// let partitioned_file = PartitionedFile::new("my_file.parquet", 1234) /// .with_extension(access_plan); /// // create a FileScanConfig to scan this file -/// let config = FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), Arc::new(ParquetSource::new(schema()))) +/// let config = FileScanConfigBuilder::new(StorageUrl::local_filesystem(), Arc::new(ParquetSource::new(schema()))) /// .with_file(partitioned_file).build(); /// // this parquet DataSourceExec will not even try to read row groups 2 and 4. Additional /// // pruning based on predicates may also happen @@ -291,6 +290,8 @@ use parquet::encryption::decrypt::FileDecryptionProperties; /// [`PhysicalExprAdapterFactory`]: datafusion_physical_expr_adapter::PhysicalExprAdapterFactory #[derive(Clone, Debug)] pub struct ParquetSource { + /// Storage resolved once when the table is constructed. + pub(crate) storage: Option>, /// Options for reading Parquet files pub(crate) table_parquet_options: TableParquetOptions, /// Optional metrics @@ -321,6 +322,103 @@ pub struct ParquetSource { } impl ParquetSource { + /// Bind file access for all executions of this source. A custom Parquet + /// reader factory, when provided, still takes precedence. + pub fn with_storage(mut self, storage: Arc) -> Self { + self.storage = Some(storage); + self + } + + fn morselizer_with_factory( + &self, + base_config: &FileScanConfig, + partition: usize, + parquet_file_reader_factory: Arc, + ) -> datafusion_common::Result> { + let expr_adapter_factory = base_config + .expr_adapter_factory + .clone() + .unwrap_or_else(|| Arc::new(DefaultPhysicalExprAdapterFactory) as _); + + #[cfg(feature = "parquet_encryption")] + let file_decryption_properties = self + .table_parquet_options() + .crypto + .file_decryption + .clone() + .map(FileDecryptionProperties::try_from) + .transpose()? + .map(Arc::new); + + let coerce_int96 = self + .table_parquet_options + .global + .coerce_int96 + .as_ref() + .map(|time_unit| parse_coerce_int96_string(time_unit.as_str()).unwrap()); + let coerce_int96_tz = self + .table_parquet_options + .global + .coerce_int96_tz + .as_ref() + .map(|tz| parse_coerce_int96_tz_string(tz)) + .transpose()?; + if coerce_int96_tz.is_some() && coerce_int96.is_none() { + warn!( + "coerce_int96_tz is set but coerce_int96 is not; the timezone will be ignored" + ); + } + + // Validate virtual columns (extension-type allowlist) and, when + // pushdown is enabled, reject predicates that reference them. Both + // checks depend only on morselizer-level state, so we pay their cost + // once per scan partition rather than per file. + // + // Gating predicate validation on `pushdown_filters` is deliberate: + // when pushdown is off the predicate stays above the scan as a + // `FilterExec` and resolves virtual columns there; the row-filter + // ban only applies to the pushdown path. + let virtual_state = build_virtual_columns_state( + self.table_schema.virtual_columns(), + self.table_schema.file_schema(), + self.predicate.as_ref(), + self.pushdown_filters(), + )?; + + Ok(Box::new(ParquetMorselizer { + partition_index: partition, + projection: self.projection.clone(), + batch_size: self + .batch_size + .expect("Batch size must set before creating ParquetMorselizer"), + limit: base_config.limit, + preserve_order: base_config.preserve_order, + predicate: self.predicate.clone(), + table_schema: self.table_schema.clone(), + metadata_size_hint: self.metadata_size_hint, + metrics: self.metrics().clone(), + parquet_file_reader_factory, + pushdown_filters: self.pushdown_filters(), + reorder_filters: self.reorder_filters(), + force_filter_selections: self.force_filter_selections(), + enable_page_index: self.enable_page_index(), + enable_bloom_filter: self.bloom_filter_on_read(), + enable_row_group_stats_pruning: self.table_parquet_options.global.pruning, + coerce_int96, + coerce_int96_tz, + #[cfg(feature = "parquet_encryption")] + file_decryption_properties, + expr_adapter_factory, + #[cfg(feature = "parquet_encryption")] + encryption_factory: self.get_encryption_factory_with_config(), + max_predicate_cache_size: self.max_predicate_cache_size(), + max_in_list_size: self.max_in_list_size(), + reverse_row_groups: self.reverse_row_groups, + sort_order_for_reorder: self.sort_order_for_reorder.clone(), + virtual_state, + })) + } + /// Create a new ParquetSource to read the data specified in the file scan /// configuration with the provided schema. /// @@ -338,6 +436,7 @@ impl ParquetSource { metrics: ExecutionPlanMetricsSet::new(), predicate: None, parquet_file_reader_factory: None, + storage: None, batch_size: None, metadata_size_hint: None, #[cfg(feature = "parquet_encryption")] @@ -562,108 +661,44 @@ impl From for Arc { impl FileSource for ParquetSource { fn create_file_opener( &self, - _object_store: Arc, + _object_store: Arc, _base_config: &FileScanConfig, _partition: usize, + _access_context: datafusion_storage::FileAccessContext, ) -> datafusion_common::Result> { datafusion_common::internal_err!( "ParquetSource::create_file_opener called but it supports the Morsel API, please use that instead" ) } - fn create_morselizer( + fn create_morselizer_with_context( &self, - object_store: Arc, base_config: &FileScanConfig, partition: usize, + context: &Arc, + access_context: datafusion_storage::FileAccessContext, ) -> datafusion_common::Result> { - let expr_adapter_factory = base_config - .expr_adapter_factory - .clone() - .unwrap_or_else(|| Arc::new(DefaultPhysicalExprAdapterFactory) as _); - - let parquet_file_reader_factory = - self.parquet_file_reader_factory.clone().unwrap_or_else(|| { - Arc::new(DefaultParquetFileReaderFactory::new(object_store)) as _ - }); - - #[cfg(feature = "parquet_encryption")] - let file_decryption_properties = self - .table_parquet_options() - .crypto - .file_decryption - .clone() - .map(FileDecryptionProperties::try_from) - .transpose()? - .map(Arc::new); - - let coerce_int96 = self - .table_parquet_options - .global - .coerce_int96 - .as_ref() - .map(|time_unit| parse_coerce_int96_string(time_unit.as_str()).unwrap()); - let coerce_int96_tz = self - .table_parquet_options - .global - .coerce_int96_tz - .as_ref() - .map(|tz| parse_coerce_int96_tz_string(tz)) - .transpose()?; - if coerce_int96_tz.is_some() && coerce_int96.is_none() { - warn!( - "coerce_int96_tz is set but coerce_int96 is not; the timezone will be ignored" - ); - } - - // Validate virtual columns (extension-type allowlist) and, when - // pushdown is enabled, reject predicates that reference them. Both - // checks depend only on morselizer-level state, so we pay their cost - // once per scan partition rather than per file. - // - // Gating predicate validation on `pushdown_filters` is deliberate: - // when pushdown is off the predicate stays above the scan as a - // `FilterExec` and resolves virtual columns there; the row-filter - // ban only applies to the pushdown path. - let virtual_state = build_virtual_columns_state( - self.table_schema.virtual_columns(), - self.table_schema.file_schema(), - self.predicate.as_ref(), - self.pushdown_filters(), - )?; - - Ok(Box::new(ParquetMorselizer { - partition_index: partition, - projection: self.projection.clone(), - batch_size: self - .batch_size - .expect("Batch size must set before creating ParquetMorselizer"), - limit: base_config.limit, - preserve_order: base_config.preserve_order, - predicate: self.predicate.clone(), - table_schema: self.table_schema.clone(), - metadata_size_hint: self.metadata_size_hint, - metrics: self.metrics().clone(), - parquet_file_reader_factory, - pushdown_filters: self.pushdown_filters(), - reorder_filters: self.reorder_filters(), - force_filter_selections: self.force_filter_selections(), - enable_page_index: self.enable_page_index(), - enable_bloom_filter: self.bloom_filter_on_read(), - enable_row_group_stats_pruning: self.table_parquet_options.global.pruning, - coerce_int96, - coerce_int96_tz, - #[cfg(feature = "parquet_encryption")] - file_decryption_properties, - expr_adapter_factory, - #[cfg(feature = "parquet_encryption")] - encryption_factory: self.get_encryption_factory_with_config(), - max_predicate_cache_size: self.max_predicate_cache_size(), - max_in_list_size: self.max_in_list_size(), - reverse_row_groups: self.reverse_row_groups, - sort_order_for_reorder: self.sort_order_for_reorder.clone(), - virtual_state, - })) + let factory = if let Some(factory) = &self.parquet_file_reader_factory { + Arc::clone(factory) + } else { + let storage = match base_config.storage.as_ref().or(self.storage.as_ref()) { + Some(storage) => Arc::clone(storage), + None => context + .runtime_env() + .storage(&base_config.object_store_url)?, + }; + Arc::new( + crate::CachedParquetFileReaderFactory::new( + storage, + context + .runtime_env() + .cache_manager + .get_file_metadata_cache(), + ) + .with_context(access_context), + ) + }; + self.morselizer_with_factory(base_config, partition, factory) } fn reorder_files( @@ -1128,12 +1163,10 @@ impl ParquetSource { node: &datafusion_proto_models::protobuf::PhysicalPlanNode, ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>, ) -> datafusion_common::Result> { - use crate::CachedParquetFileReaderFactory; use arrow::datatypes::Schema; use datafusion_common::config::TableParquetOptions; use datafusion_datasource::file_scan_config::FileScanConfig; use datafusion_datasource::source::DataSourceExec; - use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_proto_models::protobuf; let Some(protobuf::physical_plan_node::PhysicalPlanType::ParquetScan(scan)) = @@ -1199,25 +1232,8 @@ impl ParquetSource { } let table_schema = FileScanConfig::parse_table_schema_from_proto(base_conf)?; - let object_store_url = match base_conf.object_store_url.is_empty() { - false => ObjectStoreUrl::parse(&base_conf.object_store_url)?, - true => ObjectStoreUrl::local_filesystem(), - }; - let store = ctx - .task_ctx() - .runtime_env() - .object_store(object_store_url)?; - let metadata_cache = ctx - .task_ctx() - .runtime_env() - .cache_manager - .get_file_metadata_cache(); - let reader_factory = - Arc::new(CachedParquetFileReaderFactory::new(store, metadata_cache)); - - let mut source = ParquetSource::new(table_schema) - .with_parquet_file_reader_factory(reader_factory) - .with_table_parquet_options(options); + let mut source = + ParquetSource::new(table_schema).with_table_parquet_options(options); source.sort_order_for_reorder = sort_order_for_reorder; source.reverse_row_groups = scan.reverse_row_groups; diff --git a/datafusion/datasource-parquet/src/storage.rs b/datafusion/datasource-parquet/src/storage.rs new file mode 100644 index 0000000000000..542ec61aea709 --- /dev/null +++ b/datafusion/datasource-parquet/src/storage.rs @@ -0,0 +1,156 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Parquet access through a bound storage backend. +//! +//! Metadata inference and scans use the same Storage and metadata decoder. +//! The table owns its binding; replacing a runtime registration affects only +//! subsequently constructed tables. No ObjectStore client is used by this path. +use crate::ParquetFileReaderFactory; +use crate::source::ParquetSource; +use arrow::datatypes::{Schema, SchemaRef}; +use async_trait::async_trait; +use datafusion_common::{DFSchema, DataFusionError, Result}; +use datafusion_datasource::{ + PartitionedFile, file_scan_config::FileScanConfigBuilder, source::DataSourceExec, +}; +use datafusion_expr::{Expr, TableProviderFilterPushDown, TableType}; +use datafusion_physical_plan::{ExecutionPlan, metrics::ExecutionPlanMetricsSet}; +use datafusion_session::{Session, TableProvider}; +use datafusion_storage::{FileAccessContext, FileInfo, StorageUrl}; +use parquet::arrow::parquet_to_arrow_schema; +use std::sync::Arc; + +/// A fixed set of Parquet files and the backend that discovered them. +/// +/// Manifest callers can supply file metadata directly and omit listing. Files +/// discovered from a URL are snapshotted at construction; construct a new table +/// to discover additional files. Schema inference reads through the same binding +/// used by every subsequent scan. +#[derive(Debug)] +pub struct StorageParquetTable { + storage: Arc, + url: StorageUrl, + files: Vec, + schema: SchemaRef, + reader_factory: Option>, +} +impl StorageParquetTable { + /// Build a table from explicit file metadata. An explicit schema avoids + /// metadata I/O during construction and permits an empty file list. + pub async fn try_new( + storage: Arc, + files: Vec, + schema: Option, + ) -> Result { + let url = storage.url().clone(); + let schema = match schema { + Some(schema) => schema, + None => { + if files.is_empty() { + return Err(DataFusionError::Plan( + "cannot infer schema from an empty file list".into(), + )); + } + let context = FileAccessContext::new("parquet-schema-inference"); + let _cancel_on_drop = context.cancellation.clone().drop_guard(); + let mut schemas = Vec::with_capacity(files.len()); + for file in &files { + let factory = + crate::DefaultParquetFileReaderFactory::new(Arc::clone(&storage)) + .with_context(context.clone()); + let mut reader = factory.create_reader( + 0, + PartitionedFile::new_from_meta(file.clone()), + None, + &ExecutionPlanMetricsSet::new(), + )?; + let metadata = reader.get_metadata(None).await?; + schemas.push(parquet_to_arrow_schema( + metadata.file_metadata().schema_descr(), + metadata.file_metadata().key_value_metadata(), + )?); + } + Arc::new(Schema::try_merge(schemas)?) + } + }; + Ok(Self { + storage, + url, + files, + schema, + reader_factory: None, + }) + } + + /// Override the format-level reader. Inference can be avoided by supplying + /// a schema to `try_new`; the scan never replaces this custom factory. + pub fn with_parquet_file_reader_factory( + mut self, + factory: Arc, + ) -> Self { + self.reader_factory = Some(factory); + self + } +} +#[async_trait] +impl TableProvider for StorageParquetTable { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + fn table_type(&self) -> TableType { + TableType::Base + } + fn supports_filters_pushdown( + &self, + filters: &[&Expr], + ) -> Result> { + Ok(vec![TableProviderFilterPushDown::Inexact; filters.len()]) + } + async fn scan( + &self, + session: &dyn Session, + projection: Option<&[usize]>, + filters: &[Expr], + limit: Option, + ) -> Result> { + let mut source = ParquetSource::new(Arc::clone(&self.schema)) + .with_storage(Arc::clone(&self.storage)) + .with_pushdown_filters(true); + if let Some(factory) = &self.reader_factory { + source = source.with_parquet_file_reader_factory(Arc::clone(factory)); + } + if let Some(predicate) = + datafusion_expr::utils::conjunction(filters.iter().cloned()) + { + source = source.with_predicate(session.create_physical_expr( + predicate, + &DFSchema::try_from(Arc::clone(&self.schema))?, + )?); + } + let mut builder = FileScanConfigBuilder::new(self.url.clone(), Arc::new(source)); + for file in &self.files { + let partition = PartitionedFile::new_from_meta(file.clone()); + builder = builder.with_file(partition); + } + let config = builder + .with_projection_indices(projection.map(<[usize]>::to_vec))? + .with_limit(if filters.is_empty() { limit } else { None }) + .build(); + Ok(DataSourceExec::from_data_source(config)) + } +} diff --git a/datafusion/datasource-parquet/src/writer.rs b/datafusion/datasource-parquet/src/writer.rs index d37b6e26a7536..264c457b8c96e 100644 --- a/datafusion/datasource-parquet/src/writer.rs +++ b/datafusion/datasource-parquet/src/writer.rs @@ -20,9 +20,10 @@ use datafusion_common_runtime::JoinSet; use datafusion_datasource::ListingTableUrl; use datafusion_execution::TaskContext; use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties}; +use datafusion_storage::WriterOptions; use futures::StreamExt; -use object_store::buffered::BufWriter; -use object_store::path::Path; + +use datafusion_storage::path::Path; use parquet::arrow::AsyncArrowWriter; use parquet::file::properties::WriterProperties; use std::sync::Arc; @@ -36,8 +37,8 @@ pub async fn plan_to_parquet( ) -> datafusion_common::Result<()> { let path = path.as_ref(); let parsed = ListingTableUrl::parse(path)?; - let object_store_url = parsed.object_store(); - let store = task_ctx.runtime_env().object_store(&object_store_url)?; + let object_store_url = parsed.storage_url(); + let store = task_ctx.runtime_env().storage(&object_store_url)?; let mut join_set = JoinSet::new(); for i in 0..plan.output_partitioning().partition_count() { let plan: Arc = Arc::clone(&plan); @@ -46,15 +47,21 @@ pub async fn plan_to_parquet( let propclone = writer_properties.clone(); let storeref = Arc::clone(&store); - let buf_writer = BufWriter::with_capacity( - storeref, - file.clone(), - task_ctx - .session_config() - .options() - .execution - .objectstore_writer_buffer_size, - ); + let buf_writer = storeref + .writer( + &file, + WriterOptions { + buffer_size: Some( + task_ctx + .session_config() + .options() + .execution + .objectstore_writer_buffer_size, + ), + }, + datafusion_storage::FileAccessContext::new("write"), + ) + .await?; let mut stream = plan.execute(i, Arc::clone(&task_ctx))?; join_set.spawn(async move { let mut writer = diff --git a/datafusion/datasource/Cargo.toml b/datafusion/datasource/Cargo.toml index f09447b694f52..0e8d1d57ca2ed 100644 --- a/datafusion/datasource/Cargo.toml +++ b/datafusion/datasource/Cargo.toml @@ -39,6 +39,7 @@ default = ["compression"] # consumers that never serialize plans pay nothing. Mirrors the `proto` feature # on `datafusion-physical-plan`. proto = [ + "dep:datafusion-proto-models", "datafusion-physical-plan/proto", ] @@ -56,7 +57,7 @@ async-trait = { workspace = true } bytes = { workspace = true } bzip2 = { workspace = true, optional = true } chrono = { workspace = true } -datafusion-common = { workspace = true, features = ["object_store"] } +datafusion-common = { workspace = true } datafusion-common-runtime = { workspace = true } datafusion-execution = { workspace = true } datafusion-expr = { workspace = true } @@ -66,13 +67,13 @@ datafusion-physical-expr-common = { workspace = true } datafusion-physical-plan = { workspace = true } datafusion-proto-models = { workspace = true, optional = true } datafusion-session = { workspace = true } +datafusion-storage = { workspace = true } flate2 = { workspace = true, optional = true } futures = { workspace = true } glob = { workspace = true } itertools = { workspace = true } liblzma = { workspace = true, optional = true } log = { workspace = true } -object_store = { workspace = true } parking_lot = { workspace = true } rand = { workspace = true } tempfile = { workspace = true, optional = true } @@ -85,7 +86,9 @@ zstd = { workspace = true, optional = true } criterion = { workspace = true } datafusion-functions = { workspace = true } insta = { workspace = true } +object_store = { workspace = true } tempfile = { workspace = true } +test-utils = { path = "../../test-utils" } # Note: add additional linter rules in lib.rs. # Rust does not support workspace + new linter rules in subcrates yet diff --git a/datafusion/datasource/src/boundary_stream.rs b/datafusion/datasource/src/boundary_stream.rs index a1d10651b5e46..8423435595bcb 100644 --- a/datafusion/datasource/src/boundary_stream.rs +++ b/datafusion/datasource/src/boundary_stream.rs @@ -26,9 +26,9 @@ use std::sync::Arc; use std::task::{Context, Poll}; use bytes::Bytes; +use datafusion_storage::FileReader; +use futures::StreamExt; use futures::stream::{BoxStream, Stream}; -use futures::{StreamExt, TryFutureExt}; -use object_store::{GetOptions, GetRange, GetResultPayload, ObjectStore}; /// How far past `raw_end` the initial bounded fetch covers. If the terminating /// newline is not found within this window, `ScanningLastTerminator` issues @@ -62,7 +62,7 @@ enum Phase { /// `get_opts` calls (`END_SCAN_LOOKAHEAD` bytes each) until the newline is /// found or EOF is reached. pub struct AlignedBoundaryStream { - inner: BoxStream<'static, object_store::Result>, + inner: BoxStream<'static, datafusion_storage::Result>, terminator: u8, /// Effective end boundary. Set to `u64::MAX` when `end >= file_size` /// (last partition), so `FetchingChunks` never transitions to @@ -77,68 +77,11 @@ pub struct AlignedBoundaryStream { /// end-boundary processing. Consumed by `FetchingChunks` before polling /// `inner`. pending: Option, - store: Arc, - location: object_store::path::Path, + reader: Arc, /// Total file size; overflow stops when `abs_pos() >= file_size`. file_size: u64, } -/// Fetch a bounded byte range from `store` and return it as a stream -async fn get_stream( - store: Arc, - location: object_store::path::Path, - range: std::ops::Range, -) -> object_store::Result>> { - let opts = GetOptions { - range: Some(GetRange::Bounded(range.clone())), - ..Default::default() - }; - let result = store.get_opts(&location, opts).await?; - - #[cfg(not(target_arch = "wasm32"))] - if let GetResultPayload::File(mut file, _path) = result.payload { - use std::io::{Read, Seek, SeekFrom}; - const CHUNK_SIZE: u64 = 8 * 1024; - - file.seek(SeekFrom::Start(range.start)).map_err(|e| { - object_store::Error::Generic { - store: "local", - source: Box::new(e), - } - })?; - - return Ok(futures::stream::try_unfold( - (file, range.end - range.start), - move |(mut file, remaining)| async move { - if remaining == 0 { - return Ok(None); - } - let to_read = remaining.min(CHUNK_SIZE); - let cap = usize::try_from(to_read).map_err(|e| { - object_store::Error::Generic { - store: "local", - source: Box::new(e), - } - })?; - - let mut buf = Vec::with_capacity(cap); - let read = - (&mut file) - .take(to_read) - .read_to_end(&mut buf) - .map_err(|e| object_store::Error::Generic { - store: "local", - source: Box::new(e), - })?; - Ok(Some((Bytes::from(buf), (file, remaining - read as u64)))) - }, - ) - .boxed()); - } - - Ok(result.into_stream()) -} - impl AlignedBoundaryStream { /// Open a ranged byte stream from `store` and return a ready-to-poll /// `AlignedBoundaryStream`. @@ -148,16 +91,15 @@ impl AlignedBoundaryStream { /// newline is not found within that window, `ScanningLastTerminator` /// automatically issues additional `END_SCAN_LOOKAHEAD`-sized GETs /// via `store` until the newline is found or EOF is reached. - pub async fn new( - store: Arc, - location: object_store::path::Path, + pub fn new( + reader: Arc, raw_start: u64, raw_end: u64, file_size: u64, terminator: u8, - ) -> object_store::Result { + ) -> Self { if raw_start >= raw_end || raw_start >= file_size { - return Ok(Self { + return Self { inner: futures::stream::empty().boxed(), terminator, end: 0, @@ -165,10 +107,9 @@ impl AlignedBoundaryStream { fetch_start: 0, phase: Phase::Done, pending: None, - store, - location, + reader, file_size, - }); + }; } let (fetch_start, phase) = if raw_start == 0 { @@ -179,12 +120,8 @@ impl AlignedBoundaryStream { let initial_fetch_end = raw_end.saturating_add(END_SCAN_LOOKAHEAD).min(file_size); - let inner = get_stream( - Arc::clone(&store), - location.clone(), - fetch_start..initial_fetch_end, - ) - .await?; + let inner = + Arc::clone(&reader).stream(Some((fetch_start..initial_fetch_end).into())); // Last partition reads until EOF is reached — no end-boundary scanning needed. let end = if raw_end >= file_size { @@ -193,7 +130,7 @@ impl AlignedBoundaryStream { raw_end }; - Ok(Self { + Self { inner, terminator, end, @@ -201,10 +138,9 @@ impl AlignedBoundaryStream { fetch_start, phase, pending: None, - store, - location, + reader, file_size, - }) + } } /// Current absolute position in the file. @@ -214,7 +150,7 @@ impl AlignedBoundaryStream { } impl Stream for AlignedBoundaryStream { - type Item = object_store::Result; + type Item = datafusion_storage::Result; fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let this = self.get_mut(); @@ -361,11 +297,8 @@ impl Stream for AlignedBoundaryStream { let fetch_end = pos .saturating_add(END_SCAN_LOOKAHEAD) .min(this.file_size); - let store = Arc::clone(&this.store); - let location = this.location.clone(); - this.inner = get_stream(store, location, pos..fetch_end) - .try_flatten_stream() - .boxed(); + this.inner = Arc::clone(&this.reader) + .stream(Some((pos..fetch_end).into())); continue; } this.phase = Phase::Done; @@ -396,7 +329,7 @@ impl Stream for AlignedBoundaryStream { #[cfg(test)] mod tests { use super::*; - use crate::test_util::{CHUNK_SIZES, make_chunked_store}; + use crate::test_util::{CHUNK_SIZES, make_chunked_reader}; use futures::TryStreamExt; async fn collect_stream(stream: AlignedBoundaryStream) -> Vec { @@ -408,10 +341,8 @@ mod tests { // start=0, end >= file_size → pass through everything static DATA: &[u8] = b"line1\nline2\nline3\n"; for &cs in CHUNK_SIZES { - let (store, path) = make_chunked_store(DATA, cs).await; - let s = AlignedBoundaryStream::new(store, path, 0, 100, 18, b'\n') - .await - .unwrap(); + let reader = make_chunked_reader(DATA, cs); + let s = AlignedBoundaryStream::new(reader, 0, 100, 18, b'\n'); assert_eq!(collect_stream(s).await, DATA, "chunk_size={cs}"); } } @@ -424,10 +355,8 @@ mod tests { // Should skip the leading '\n' and yield "line2\nline3\n". static DATA: &[u8] = b"line1\nline2\nline3\n"; for &cs in CHUNK_SIZES { - let (store, path) = make_chunked_store(DATA, cs).await; - let s = AlignedBoundaryStream::new(store, path, 6, 100, 18, b'\n') - .await - .unwrap(); + let reader = make_chunked_reader(DATA, cs); + let s = AlignedBoundaryStream::new(reader, 6, 100, 18, b'\n'); assert_eq!( collect_stream(s).await, b"line2\nline3\n", @@ -442,10 +371,8 @@ mod tests { // Should skip "ne1\n" and yield "line2\nline3\n". static DATA: &[u8] = b"line1\nline2\nline3\n"; for &cs in CHUNK_SIZES { - let (store, path) = make_chunked_store(DATA, cs).await; - let s = AlignedBoundaryStream::new(store, path, 3, 100, 18, b'\n') - .await - .unwrap(); + let reader = make_chunked_reader(DATA, cs); + let s = AlignedBoundaryStream::new(reader, 3, 100, 18, b'\n'); assert_eq!( collect_stream(s).await, b"line2\nline3\n", @@ -462,10 +389,8 @@ mod tests { // Should yield "line1\nline2\n" (continue past end to find newline). static DATA: &[u8] = b"line1\nline2\nline3\n"; for &cs in CHUNK_SIZES { - let (store, path) = make_chunked_store(DATA, cs).await; - let s = AlignedBoundaryStream::new(store, path, 0, 8, 18, b'\n') - .await - .unwrap(); + let reader = make_chunked_reader(DATA, cs); + let s = AlignedBoundaryStream::new(reader, 0, 8, 18, b'\n'); assert_eq!( collect_stream(s).await, b"line1\nline2\n", @@ -479,10 +404,8 @@ mod tests { // end >= file_size → no end scanning, pass through everything. static DATA: &[u8] = b"line1\nline2\n"; for &cs in CHUNK_SIZES { - let (store, path) = make_chunked_store(DATA, cs).await; - let s = AlignedBoundaryStream::new(store, path, 0, 12, 12, b'\n') - .await - .unwrap(); + let reader = make_chunked_reader(DATA, cs); + let s = AlignedBoundaryStream::new(reader, 0, 12, 12, b'\n'); assert_eq!(collect_stream(s).await, DATA, "chunk_size={cs}"); } } @@ -493,10 +416,8 @@ mod tests { // No complete line → empty output. static DATA: &[u8] = b"abcdef"; for &cs in CHUNK_SIZES { - let (store, path) = make_chunked_store(DATA, cs).await; - let s = AlignedBoundaryStream::new(store, path, 2, 6, 6, b'\n') - .await - .unwrap(); + let reader = make_chunked_reader(DATA, cs); + let s = AlignedBoundaryStream::new(reader, 2, 6, 6, b'\n'); assert!(collect_stream(s).await.is_empty(), "chunk_size={cs}"); } } @@ -511,10 +432,8 @@ mod tests { // Expected: "line2\nline3\n" static DATA: &[u8] = b"line1\nline2\nline3\nline4\n"; for &cs in CHUNK_SIZES { - let (store, path) = make_chunked_store(DATA, cs).await; - let s = AlignedBoundaryStream::new(store, path, 3, 14, 24, b'\n') - .await - .unwrap(); + let reader = make_chunked_reader(DATA, cs); + let s = AlignedBoundaryStream::new(reader, 3, 14, 24, b'\n'); assert_eq!( collect_stream(s).await, b"line2\nline3\n", @@ -531,10 +450,8 @@ mod tests { // start=0, end=7 (mid "line2"), file_size=18 → "line1\nline2\n" static DATA: &[u8] = b"line1\nline2\nline3\n"; for &cs in CHUNK_SIZES { - let (store, path) = make_chunked_store(DATA, cs).await; - let s = AlignedBoundaryStream::new(store, path, 0, 7, 18, b'\n') - .await - .unwrap(); + let reader = make_chunked_reader(DATA, cs); + let s = AlignedBoundaryStream::new(reader, 0, 7, 18, b'\n'); assert_eq!( collect_stream(s).await, b"line1\nline2\n", @@ -548,51 +465,24 @@ mod tests { // start >= end — no complete line can exist, regardless of data. static DATA: &[u8] = b"line1\nline2\n"; for &cs in CHUNK_SIZES { - let (store, path) = make_chunked_store(DATA, cs).await; + let reader = make_chunked_reader(DATA, cs); // start > end (non-zero start) - let s = AlignedBoundaryStream::new( - Arc::clone(&store), - path.clone(), - 10, - 5, - 20, - b'\n', - ) - .await - .unwrap(); + let s = AlignedBoundaryStream::new(Arc::clone(&reader), 10, 5, 20, b'\n'); assert!( collect_stream(s).await.is_empty(), "start>end chunk_size={cs}" ); // start == end == 0 (zero start, previously unguarded) - let s = AlignedBoundaryStream::new( - Arc::clone(&store), - path.clone(), - 0, - 0, - 12, - b'\n', - ) - .await - .unwrap(); + let s = AlignedBoundaryStream::new(Arc::clone(&reader), 0, 0, 12, b'\n'); assert!( collect_stream(s).await.is_empty(), "start==end==0 chunk_size={cs}" ); // start == end (non-zero) - let s = AlignedBoundaryStream::new( - Arc::clone(&store), - path.clone(), - 6, - 6, - 12, - b'\n', - ) - .await - .unwrap(); + let s = AlignedBoundaryStream::new(Arc::clone(&reader), 6, 6, 12, b'\n'); assert!( collect_stream(s).await.is_empty(), "start==end==6 chunk_size={cs}" @@ -607,10 +497,8 @@ mod tests { // Start aligns past "abcdef\n", yielding "line2\n". static DATA: &[u8] = b"abcdef\nline2\n"; for &cs in CHUNK_SIZES { - let (store, path) = make_chunked_store(DATA, cs).await; - let s = AlignedBoundaryStream::new(store, path, 1, 100, 13, b'\n') - .await - .unwrap(); + let reader = make_chunked_reader(DATA, cs); + let s = AlignedBoundaryStream::new(reader, 1, 100, 13, b'\n'); assert_eq!(collect_stream(s).await, b"line2\n", "chunk_size={cs}"); } } @@ -623,10 +511,8 @@ mod tests { // start=0, end=6 → byte 5 is '\n' → yield only "line1\n". static DATA: &[u8] = b"line1\nline2\nline3\n"; for &cs in CHUNK_SIZES { - let (store, path) = make_chunked_store(DATA, cs).await; - let s = AlignedBoundaryStream::new(store, path, 0, 6, 18, b'\n') - .await - .unwrap(); + let reader = make_chunked_reader(DATA, cs); + let s = AlignedBoundaryStream::new(reader, 0, 6, 18, b'\n'); assert_eq!(collect_stream(s).await, b"line1\n", "chunk_size={cs}"); } } @@ -640,45 +526,30 @@ mod tests { static DATA: &[u8] = b"line1\nline2\nline3\n"; // 18 bytes for &cs in CHUNK_SIZES { - let (store, path) = make_chunked_store(DATA, cs).await; - let r1 = collect_stream( - AlignedBoundaryStream::new( - Arc::clone(&store), - path.clone(), - 0, - 6, - 18, - b'\n', - ) - .await - .unwrap(), - ) + let reader = make_chunked_reader(DATA, cs); + let r1 = collect_stream(AlignedBoundaryStream::new( + Arc::clone(&reader), + 0, + 6, + 18, + b'\n', + )) .await; - let r2 = collect_stream( - AlignedBoundaryStream::new( - Arc::clone(&store), - path.clone(), - 6, - 12, - 18, - b'\n', - ) - .await - .unwrap(), - ) + let r2 = collect_stream(AlignedBoundaryStream::new( + Arc::clone(&reader), + 6, + 12, + 18, + b'\n', + )) .await; - let r3 = collect_stream( - AlignedBoundaryStream::new( - Arc::clone(&store), - path.clone(), - 12, - 18, - 18, - b'\n', - ) - .await - .unwrap(), - ) + let r3 = collect_stream(AlignedBoundaryStream::new( + Arc::clone(&reader), + 12, + 18, + 18, + b'\n', + )) .await; assert_eq!(r1, b"line1\n", "p1 chunk_size={cs}"); @@ -702,10 +573,8 @@ mod tests { // aligned start = 11, which is >= end = 6 → empty. static DATA: &[u8] = b"abcdefghij\nkl\n"; for &cs in CHUNK_SIZES { - let (store, path) = make_chunked_store(DATA, cs).await; - let s = AlignedBoundaryStream::new(store, path, 3, 6, 14, b'\n') - .await - .unwrap(); + let reader = make_chunked_reader(DATA, cs); + let s = AlignedBoundaryStream::new(reader, 3, 6, 14, b'\n'); assert!(collect_stream(s).await.is_empty(), "chunk_size={cs}"); } } @@ -719,55 +588,40 @@ mod tests { static DATA: &[u8] = b"aaa\nbbb\nccc\n"; // 12 bytes for &cs in CHUNK_SIZES { - let (store, path) = make_chunked_store(DATA, cs).await; + let reader = make_chunked_reader(DATA, cs); // [0, 5): no start alignment; end=5 mid "bbb", scans to '\n' at 7. - let r1 = collect_stream( - AlignedBoundaryStream::new( - Arc::clone(&store), - path.clone(), - 0, - 5, - 12, - b'\n', - ) - .await - .unwrap(), - ) + let r1 = collect_stream(AlignedBoundaryStream::new( + Arc::clone(&reader), + 0, + 5, + 12, + b'\n', + )) .await; // [5, 10): fetch_start=4, bytes from offset 4: "bbb\nccc\n". // '\n' at pos 3 → aligned start=8 ("ccc\n"). End=10 mid "ccc", // scans to '\n' at 11 → yields "ccc\n". - let r2 = collect_stream( - AlignedBoundaryStream::new( - Arc::clone(&store), - path.clone(), - 5, - 10, - 12, - b'\n', - ) - .await - .unwrap(), - ) + let r2 = collect_stream(AlignedBoundaryStream::new( + Arc::clone(&reader), + 5, + 10, + 12, + b'\n', + )) .await; // [10, 12): fetch_start=9, bytes from offset 9: "cc\n". // '\n' at pos 2 → aligned start=12. end=12==file_size → end=MAX. // Remainder after '\n' is empty; Passthrough polls inner → Done. - let r3 = collect_stream( - AlignedBoundaryStream::new( - Arc::clone(&store), - path.clone(), - 10, - 12, - 12, - b'\n', - ) - .await - .unwrap(), - ) + let r3 = collect_stream(AlignedBoundaryStream::new( + Arc::clone(&reader), + 10, + 12, + 12, + b'\n', + )) .await; assert_eq!(r1, b"aaa\nbbb\n", "p1 chunk_size={cs}"); @@ -788,34 +642,16 @@ mod tests { // until EOF is reached and yields the final incomplete line as-is. static DATA: &[u8] = b"line1\nline2"; // 11 bytes, no trailing '\n' for &cs in CHUNK_SIZES { - let (store, path) = make_chunked_store(DATA, cs).await; + let reader = make_chunked_reader(DATA, cs); // Single partition covering the whole file. - let s = AlignedBoundaryStream::new( - Arc::clone(&store), - path.clone(), - 0, - 11, - 11, - b'\n', - ) - .await - .unwrap(); + let s = AlignedBoundaryStream::new(Arc::clone(&reader), 0, 11, 11, b'\n'); assert_eq!(collect_stream(s).await, DATA, "chunk_size={cs}"); // Last partition starting mid-file (start=6, fetch_start=5). // Bytes from offset 5: "\nline2". // StartAlign consumes '\n', remainder "line2" is yielded as-is. - let s = AlignedBoundaryStream::new( - Arc::clone(&store), - path.clone(), - 6, - 11, - 11, - b'\n', - ) - .await - .unwrap(); + let s = AlignedBoundaryStream::new(Arc::clone(&reader), 6, 11, 11, b'\n'); assert_eq!(collect_stream(s).await, b"line2", "tail chunk_size={cs}"); } } @@ -841,34 +677,24 @@ mod tests { let file_size = data.len() as u64; for &cs in CHUNK_SIZES { - let (store, path) = make_chunked_store(&data, cs).await; - - let r1 = collect_stream( - AlignedBoundaryStream::new( - Arc::clone(&store), - path.clone(), - 0, - 1, - file_size, - b'\n', - ) - .await - .unwrap(), - ) + let reader = make_chunked_reader(&data, cs); + + let r1 = collect_stream(AlignedBoundaryStream::new( + Arc::clone(&reader), + 0, + 1, + file_size, + b'\n', + )) .await; - let r2 = collect_stream( - AlignedBoundaryStream::new( - Arc::clone(&store), - path.clone(), - 1, - file_size, - file_size, - b'\n', - ) - .await - .unwrap(), - ) + let r2 = collect_stream(AlignedBoundaryStream::new( + Arc::clone(&reader), + 1, + file_size, + file_size, + b'\n', + )) .await; assert_eq!(r1, long_line, "p1 chunk_size={cs}"); diff --git a/datafusion/datasource/src/display.rs b/datafusion/datasource/src/display.rs index 0f59e33ff9eac..eb99047fcfbef 100644 --- a/datafusion/datasource/src/display.rs +++ b/datafusion/datasource/src/display.rs @@ -135,7 +135,7 @@ mod tests { use super::*; use datafusion_physical_plan::{DefaultDisplay, VerboseDisplay}; - use object_store::{ObjectMeta, path::Path}; + use datafusion_storage::{FileInfo as ObjectMeta, path::Path}; use crate::PartitionedFile; use chrono::Utc; diff --git a/datafusion/datasource/src/file.rs b/datafusion/datasource/src/file.rs index f1a94f2e12363..8aabe627ae956 100644 --- a/datafusion/datasource/src/file.rs +++ b/datafusion/datasource/src/file.rs @@ -25,7 +25,8 @@ use std::sync::Arc; use crate::file_groups::FileGroupPartitioner; use crate::file_scan_config::FileScanConfig; use crate::file_stream::FileOpener; -use crate::morsel::{FileOpenerMorselizer, Morselizer}; +use crate::morsel::FileOpenerMorselizer; +use crate::morsel::Morselizer; #[expect(deprecated)] use crate::schema_adapter::SchemaAdapterFactory; use datafusion_common::config::ConfigOptions; @@ -39,7 +40,7 @@ use datafusion_physical_plan::filter_pushdown::{FilterPushdownPropagation, Pushe use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; -use object_store::ObjectStore; +use datafusion_storage::StorageBinding; /// Helper function to convert any type implementing [`FileSource`] to `Arc` pub fn as_file_source(source: T) -> Arc { @@ -64,15 +65,33 @@ pub fn as_file_source(source: T) -> Arc /// /// [`DataSource`]: crate::source::DataSource pub trait FileSource: Any + Send + Sync { + /// Create a morselizer using the plan's immutable storage binding. + fn create_morselizer_with_context( + &self, + base_config: &FileScanConfig, + partition: usize, + context: &Arc, + access_context: datafusion_storage::FileAccessContext, + ) -> Result> { + let storage = match &base_config.storage { + Some(storage) => Arc::clone(storage), + None => context + .runtime_env() + .storage(&base_config.object_store_url)?, + }; + self.create_morselizer(storage, base_config, partition, access_context) + } + /// Creates a `dyn FileOpener` based on given parameters. /// /// Note: File sources with a native morsel implementation should return an /// error from this method and implementing [`Self::create_morselizer`] instead. fn create_file_opener( &self, - object_store: Arc, + object_store: Arc, base_config: &FileScanConfig, partition: usize, + access_context: datafusion_storage::FileAccessContext, ) -> Result>; /// Creates a `dyn Morselizer` based on given parameters. @@ -84,11 +103,17 @@ pub trait FileSource: Any + Send + Sync { /// implementing this method. fn create_morselizer( &self, - object_store: Arc, + object_store: Arc, base_config: &FileScanConfig, partition: usize, + access_context: datafusion_storage::FileAccessContext, ) -> Result> { - let opener = self.create_file_opener(object_store, base_config, partition)?; + let opener = self.create_file_opener( + object_store, + base_config, + partition, + access_context, + )?; Ok(Box::new(FileOpenerMorselizer::new(opener))) } diff --git a/datafusion/datasource/src/file_compression_type.rs b/datafusion/datasource/src/file_compression_type.rs index 89efb580652b1..8709706157093 100644 --- a/datafusion/datasource/src/file_compression_type.rs +++ b/datafusion/datasource/src/file_compression_type.rs @@ -45,7 +45,6 @@ use futures::TryStreamExt; use futures::stream::BoxStream; #[cfg(feature = "compression")] use liblzma::read::XzDecoder; -use object_store::buffered::BufWriter; use tokio::io::AsyncWrite; #[cfg(feature = "compression")] use tokio_util::io::{ReaderStream, StreamReader}; @@ -158,7 +157,7 @@ impl FileCompressionType { /// according to this `FileCompressionType` using the default compression level. pub fn convert_async_writer( &self, - w: BufWriter, + w: impl AsyncWrite + Send + Unpin + 'static, ) -> Result> { self.convert_async_writer_with_level(w, None) } @@ -170,7 +169,7 @@ impl FileCompressionType { /// compression level. If `None`, the default level for each algorithm is used. pub fn convert_async_writer_with_level( &self, - w: BufWriter, + w: impl AsyncWrite + Send + Unpin + 'static, compression_level: Option, ) -> Result> { #[cfg(feature = "compression")] diff --git a/datafusion/datasource/src/file_format.rs b/datafusion/datasource/src/file_format.rs index 2c2a94e8cc4b4..76e1e337888b5 100644 --- a/datafusion/datasource/src/file_format.rs +++ b/datafusion/datasource/src/file_format.rs @@ -39,7 +39,7 @@ use datafusion_physical_plan::ExecutionPlan; use datafusion_session::Session; use async_trait::async_trait; -use object_store::{ObjectMeta, ObjectStore}; +use datafusion_storage::{FileInfo, StorageBinding}; /// Default max records to scan to infer the schema pub const DEFAULT_SCHEMA_INFER_MAX_RECORD: usize = 1000; @@ -120,8 +120,8 @@ pub trait FileFormat: Any + Send + Sync + fmt::Debug { async fn infer_schema( &self, state: &dyn Session, - store: &Arc, - objects: &[ObjectMeta], + store: &Arc, + objects: &[FileInfo], ) -> Result; /// Infer the statistics for the provided object. The cost and accuracy of the @@ -134,9 +134,9 @@ pub trait FileFormat: Any + Send + Sync + fmt::Debug { async fn infer_stats( &self, state: &dyn Session, - store: &Arc, + store: &Arc, table_schema: SchemaRef, - object: &ObjectMeta, + object: &FileInfo, ) -> Result; /// Infer the ordering (sort order) for the provided object from file metadata. @@ -151,9 +151,9 @@ pub trait FileFormat: Any + Send + Sync + fmt::Debug { async fn infer_ordering( &self, _state: &dyn Session, - _store: &Arc, + _store: &Arc, _table_schema: SchemaRef, - _object: &ObjectMeta, + _object: &FileInfo, ) -> Result> { Ok(None) } @@ -169,9 +169,9 @@ pub trait FileFormat: Any + Send + Sync + fmt::Debug { async fn infer_stats_and_ordering( &self, state: &dyn Session, - store: &Arc, + store: &Arc, table_schema: SchemaRef, - object: &ObjectMeta, + object: &FileInfo, ) -> Result { let statistics = self .infer_stats(state, store, Arc::clone(&table_schema), object) diff --git a/datafusion/datasource/src/file_scan_config/mod.rs b/datafusion/datasource/src/file_scan_config/mod.rs index 4e72b5e83bd23..2da367adf3b63 100644 --- a/datafusion/datasource/src/file_scan_config/mod.rs +++ b/datafusion/datasource/src/file_scan_config/mod.rs @@ -42,10 +42,9 @@ use datafusion_common::{ Constraint, Constraints, Result, ScalarValue, Statistics, internal_datafusion_err, internal_err, }; -use datafusion_execution::{ - SendableRecordBatchStream, TaskContext, object_store::ObjectStoreUrl, -}; +use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_expr::Operator; +use datafusion_storage::StorageUrl; use crate::source::OpenArgs; use datafusion_common::stats::{Precision, is_known_empty}; @@ -89,7 +88,7 @@ use std::{fmt::Debug, fmt::Formatter, fmt::Result as FmtResult, sync::Arc}; /// ``` /// # use std::sync::Arc; /// # use arrow::datatypes::{Field, Fields, DataType, Schema, SchemaRef}; -/// # use object_store::ObjectStore; +/// # use datafusion_storage::StorageBinding; /// # use datafusion_common::Result; /// # use datafusion_common::tree_node::TreeNodeRecursion; /// # use datafusion_datasource::file::FileSource; @@ -100,7 +99,7 @@ use std::{fmt::Debug, fmt::Formatter, fmt::Result as FmtResult, sync::Arc}; /// # use datafusion_datasource::file_stream::FileOpener; /// # use datafusion_datasource::source::DataSourceExec; /// # use datafusion_datasource::table_schema::TableSchema; -/// # use datafusion_execution::object_store::ObjectStoreUrl; +/// # use datafusion_storage::StorageUrl; /// # use datafusion_physical_expr::projection::ProjectionExprs; /// # use datafusion_physical_plan::ExecutionPlan; /// # use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; @@ -116,7 +115,7 @@ use std::{fmt::Debug, fmt::Formatter, fmt::Result as FmtResult, sync::Arc}; /// # table_schema: TableSchema, /// # }; /// # impl FileSource for ParquetSource { -/// # fn create_file_opener(&self, _: Arc, _: &FileScanConfig, _: usize) -> Result> { unimplemented!() } +/// # fn create_file_opener(&self, _: Arc, _: &FileScanConfig, _: usize, _: datafusion_storage::FileAccessContext) -> Result> { unimplemented!() } /// # fn table_schema(&self) -> &TableSchema { &self.table_schema } /// # fn with_batch_size(&self, _: usize) -> Arc { unimplemented!() } /// # fn metrics(&self) -> &ExecutionPlanMetricsSet { unimplemented!() } @@ -129,7 +128,7 @@ use std::{fmt::Debug, fmt::Formatter, fmt::Result as FmtResult, sync::Arc}; /// # fn new(table_schema: impl Into) -> Self { Self {table_schema: table_schema.into()} } /// # } /// // create FileScan config for reading parquet files from file:// -/// let object_store_url = ObjectStoreUrl::local_filesystem(); +/// let object_store_url = StorageUrl::local_filesystem(); /// let file_source = Arc::new(ParquetSource::new(file_schema.clone())); /// let config = FileScanConfigBuilder::new(object_store_url, file_source) /// .with_limit(Some(1000)) // read only the first 1000 records @@ -151,18 +150,10 @@ use std::{fmt::Debug, fmt::Formatter, fmt::Result as FmtResult, sync::Arc}; /// [`DataSourceExec::from_data_source`]: crate::source::DataSourceExec::from_data_source #[derive(Clone)] pub struct FileScanConfig { - /// Object store URL, used to get an [`ObjectStore`] instance from - /// [`RuntimeEnv::object_store`] - /// - /// This `ObjectStoreUrl` should be the prefix of the absolute url for files - /// as `file://` or `s3://my_bucket`. It should not include the path to the - /// file itself. The relevant URL prefix must be registered via - /// [`RuntimeEnv::register_object_store`] - /// - /// [`ObjectStore`]: object_store::ObjectStore - /// [`RuntimeEnv::register_object_store`]: datafusion_execution::runtime_env::RuntimeEnv::register_object_store - /// [`RuntimeEnv::object_store`]: datafusion_execution::runtime_env::RuntimeEnv::object_store - pub object_store_url: ObjectStoreUrl, + /// Storage namespace for these files, such as `file://` or `s3://bucket`. + /// A bound configuration retains its storage; otherwise execution resolves + /// this namespace from the runtime storage registry. + pub object_store_url: StorageUrl, /// List of files to be processed, grouped into partitions /// /// Each file must have a schema of `file_schema` or a subset. If @@ -201,6 +192,8 @@ pub struct FileScanConfig { pub file_compression_type: FileCompressionType, /// File source such as `ParquetSource`, `CsvSource`, `JsonSource`, etc. pub file_source: Arc, + /// Storage resolved during planning, retained across execution. + pub storage: Option>, /// Batch size while creating new batches /// Defaults to [`datafusion_common::config::ExecutionOptions`] batch_size. pub batch_size: Option, @@ -238,7 +231,7 @@ pub struct FileScanConfig { /// # use datafusion_datasource::file_groups::FileGroup; /// # use datafusion_datasource::PartitionedFile; /// # use datafusion_datasource::table_schema::TableSchema; -/// # use datafusion_execution::object_store::ObjectStoreUrl; +/// # use datafusion_storage::StorageUrl; /// # use datafusion_common::Statistics; /// # use datafusion_datasource::file::FileSource; /// @@ -262,7 +255,7 @@ pub struct FileScanConfig { /// /// // Create a builder for scanning Parquet files from a local filesystem /// let config = FileScanConfigBuilder::new( -/// ObjectStoreUrl::local_filesystem(), +/// StorageUrl::local_filesystem(), /// file_source, /// ) /// // Set a limit of 1000 rows @@ -284,8 +277,9 @@ pub struct FileScanConfig { /// ``` #[derive(Clone)] pub struct FileScanConfigBuilder { - object_store_url: ObjectStoreUrl, + object_store_url: StorageUrl, file_source: Arc, + storage: Option>, limit: Option, preserve_order: bool, constraints: Option, @@ -305,13 +299,11 @@ impl FileScanConfigBuilder { /// * `object_store_url`: See [`FileScanConfig::object_store_url`] /// * `file_source`: See [`FileScanConfig::file_source`]. The file source must have /// a schema set via its constructor. - pub fn new( - object_store_url: ObjectStoreUrl, - file_source: Arc, - ) -> Self { + pub fn new(object_store_url: StorageUrl, file_source: Arc) -> Self { Self { object_store_url, file_source, + storage: None, file_groups: vec![], statistics: None, output_ordering: vec![], @@ -348,6 +340,15 @@ impl FileScanConfigBuilder { /// /// This method allows you to change the file source implementation (e.g. /// ParquetSource, CsvSource, etc.) after the builder has been created. + /// Retain the complete binding selected during planning. + pub fn with_storage( + mut self, + storage: Arc, + ) -> Self { + self.storage = Some(storage); + self + } + pub fn with_source(mut self, file_source: Arc) -> Self { self.file_source = file_source; self @@ -530,6 +531,7 @@ impl FileScanConfigBuilder { let Self { object_store_url, file_source, + storage, limit, preserve_order, constraints, @@ -556,6 +558,7 @@ impl FileScanConfigBuilder { FileScanConfig { object_store_url, file_source, + storage, limit, preserve_order, constraints, @@ -599,6 +602,7 @@ fn add_key_distinct_counts(constraints: &Constraints, statistics: &mut Statistic impl From for FileScanConfigBuilder { fn from(config: FileScanConfig) -> Self { Self { + storage: config.storage, object_store_url: config.object_store_url, file_source: Arc::::clone(&config.file_source), file_groups: config.file_groups, @@ -731,14 +735,22 @@ impl DataSource for FileScanConfig { context, sibling_state, } = args; - let object_store = context.runtime_env().object_store(&self.object_store_url)?; let batch_size = self .batch_size .unwrap_or_else(|| context.session_config().batch_size()); let source = self.file_source.with_batch_size(batch_size); - let morselizer = source.create_morselizer(object_store, self, partition)?; + let access_context = datafusion_storage::FileAccessContext::new( + context.task_id().unwrap_or_else(|| context.session_id()), + ); + let cancel_on_drop = access_context.cancellation.clone().drop_guard(); + let morselizer = source.create_morselizer_with_context( + self, + partition, + &context, + access_context, + )?; // Extract the shared work source from the sibling state if it exists. // This allows multiple sibling streams to steal work from a single @@ -754,7 +766,20 @@ impl DataSource for FileScanConfig { .with_morselizer(morselizer) .with_metrics(source.metrics()) .build()?; - Ok(Box::pin(cooperative(stream))) + use datafusion_execution::RecordBatchStream; + use futures::StreamExt; + let schema = stream.schema(); + let stream = futures::stream::unfold( + (cooperative(stream), cancel_on_drop), + |(mut stream, guard)| async move { + stream.next().await.map(|batch| (batch, (stream, guard))) + }, + ); + Ok(Box::pin( + datafusion_physical_plan::stream::RecordBatchStreamAdapter::new( + schema, stream, + ), + )) } fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> FmtResult { @@ -1684,10 +1709,10 @@ mod tests { use datafusion_physical_plan::proto::{ExecutionPlanEncode, ExecutionPlanEncodeCtx}; #[cfg(feature = "proto")] use datafusion_proto_models::protobuf::{PhysicalExprNode, PhysicalPlanNode}; + use datafusion_storage::StorageBinding; use futures::FutureExt as _; use futures::StreamExt as _; use futures::stream; - use object_store::ObjectStore; use std::fmt::Debug; #[derive(Clone)] @@ -1708,9 +1733,10 @@ mod tests { impl FileSource for InexactSortPushdownSource { fn create_file_opener( &self, - _object_store: Arc, + _object_store: Arc, _base_config: &FileScanConfig, _partition: usize, + _access_context: datafusion_storage::FileAccessContext, ) -> Result> { unimplemented!() } @@ -1770,9 +1796,10 @@ mod tests { impl FileSource for ProtoHookSource { fn create_file_opener( &self, - _object_store: Arc, + _object_store: Arc, _base_config: &FileScanConfig, _partition: usize, + _access_context: datafusion_storage::FileAccessContext, ) -> Result> { internal_err!("not needed for proto delegation test") } @@ -1848,8 +1875,7 @@ mod tests { )])); let source = Arc::new(ProtoHookSource::new(TableSchema::from(&schema))); let config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source) - .build(); + FileScanConfigBuilder::new(StorageUrl::local_filesystem(), source).build(); let exec = DataSourceExec::from_data_source(config); let encoder = UnusedPlanEncoder; let ctx = ExecutionPlanEncodeCtx::new(&encoder); @@ -1893,7 +1919,7 @@ mod tests { use datafusion_common::DFSchema; use datafusion_expr::execution_props::ExecutionProps; use datafusion_expr::physical_planning_context::PhysicalPlanningContext; - use object_store::{ObjectMeta, path::Path}; + use datafusion_storage::{FileInfo as ObjectMeta, path::Path}; struct File { name: &'static str, @@ -2209,7 +2235,7 @@ mod tests { projection: Option>, ) -> FileScanConfig { FileScanConfigBuilder::new( - ObjectStoreUrl::parse("test:///").unwrap(), + StorageUrl::parse("test:///").unwrap(), Arc::new(MockSource::new(table_schema)), ) .with_statistics(statistics) @@ -2328,7 +2354,7 @@ mod tests { ) .build(); FileScanConfigBuilder::new( - ObjectStoreUrl::parse("test:///").unwrap(), + StorageUrl::parse("test:///").unwrap(), Arc::new(MockSource::new(table_schema.clone())), ) .with_projection_indices(projection) @@ -2340,7 +2366,7 @@ mod tests { #[test] fn test_file_scan_config_builder() { let file_schema = aggr_test_schema(); - let object_store_url = ObjectStoreUrl::parse("test:///").unwrap(); + let object_store_url = StorageUrl::parse("test:///").unwrap(); let table_schema = TableSchemaBuilder::from(&file_schema) .with_table_partition_cols(vec![Arc::new(Field::new( @@ -2408,7 +2434,7 @@ mod tests { #[test] fn equivalence_properties_after_schema_change() { let file_schema = aggr_test_schema(); - let object_store_url = ObjectStoreUrl::parse("test:///").unwrap(); + let object_store_url = StorageUrl::parse("test:///").unwrap(); let table_schema = TableSchema::from(&file_schema); @@ -2461,7 +2487,7 @@ mod tests { #[test] fn test_file_scan_config_builder_defaults() { let file_schema = aggr_test_schema(); - let object_store_url = ObjectStoreUrl::parse("test:///").unwrap(); + let object_store_url = StorageUrl::parse("test:///").unwrap(); let table_schema = TableSchema::from(&file_schema); @@ -2517,7 +2543,7 @@ mod tests { #[test] fn test_file_scan_config_builder_new_from() { let schema = aggr_test_schema(); - let object_store_url = ObjectStoreUrl::parse("test:///").unwrap(); + let object_store_url = StorageUrl::parse("test:///").unwrap(); let partition_cols = vec![Field::new( "date", wrap_partition_type_in_dict(DataType::Utf8), @@ -2865,7 +2891,7 @@ mod tests { // Create a FileScanConfig with projection: only keep columns 0 and 2 let config = FileScanConfigBuilder::new( - ObjectStoreUrl::parse("test:///").unwrap(), + StorageUrl::parse("test:///").unwrap(), Arc::new(MockSource::new(table_schema.clone())), ) .with_projection_indices(Some(vec![0, 2])) @@ -2935,7 +2961,7 @@ mod tests { let table_schema = TableSchema::from(&schema); let config = FileScanConfigBuilder::new( - ObjectStoreUrl::parse("test:///").unwrap(), + StorageUrl::parse("test:///").unwrap(), Arc::new(MockSource::new(table_schema.clone()).with_filter(Arc::new( Literal::new(ScalarValue::Boolean(Some(true))), ))), @@ -2967,7 +2993,8 @@ mod tests { ); let config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + FileScanConfigBuilder::new(StorageUrl::local_filesystem(), file_source) + .with_storage(test_utils::storage::local()) .with_file_group(FileGroup::new(vec![ PartitionedFile::new("file1.parquet", 100), PartitionedFile::new("file2.parquet", 100), @@ -3205,7 +3232,7 @@ mod tests { let sort_expr_asc = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0))); let config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + FileScanConfigBuilder::new(StorageUrl::local_filesystem(), file_source) .with_file_groups(file_groups) .with_output_ordering(vec![ LexOrdering::new(vec![sort_expr_asc.clone()]).unwrap(), @@ -3281,9 +3308,10 @@ mod tests { impl FileSource for ExactSortPushdownSource { fn create_file_opener( &self, - _object_store: Arc, + _object_store: Arc, _base_config: &FileScanConfig, _partition: usize, + _access_context: datafusion_storage::FileAccessContext, ) -> Result> { unimplemented!() } @@ -3337,7 +3365,7 @@ mod tests { let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0))); let config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + FileScanConfigBuilder::new(StorageUrl::local_filesystem(), file_source) .with_file_groups(file_groups) .build(); @@ -3371,7 +3399,7 @@ mod tests { let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0))); let config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + FileScanConfigBuilder::new(StorageUrl::local_filesystem(), file_source) .with_file_groups(file_groups) .build(); @@ -3401,7 +3429,7 @@ mod tests { }, ); let config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + FileScanConfigBuilder::new(StorageUrl::local_filesystem(), file_source) .with_file_groups(file_groups) .build(); @@ -3435,7 +3463,7 @@ mod tests { ])]; let config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + FileScanConfigBuilder::new(StorageUrl::local_filesystem(), file_source) .with_file_groups(file_groups) .with_output_ordering(vec![ LexOrdering::new(vec![sort_expr.clone()]).unwrap(), @@ -3469,7 +3497,7 @@ mod tests { ])]; let config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + FileScanConfigBuilder::new(StorageUrl::local_filesystem(), file_source) .with_file_groups(file_groups) .with_output_ordering(vec![ LexOrdering::new(vec![sort_expr.clone()]).unwrap(), @@ -3503,7 +3531,7 @@ mod tests { ])]; let config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + FileScanConfigBuilder::new(StorageUrl::local_filesystem(), file_source) .with_file_groups(file_groups) .with_output_ordering(vec![ LexOrdering::new(vec![sort_expr.clone()]).unwrap(), @@ -3539,7 +3567,7 @@ mod tests { let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0))); let config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + FileScanConfigBuilder::new(StorageUrl::local_filesystem(), file_source) .with_file_groups(file_groups) .build(); @@ -3571,7 +3599,7 @@ mod tests { let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0))); let config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + FileScanConfigBuilder::new(StorageUrl::local_filesystem(), file_source) .with_file_groups(file_groups) .build(); @@ -3611,7 +3639,7 @@ mod tests { let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0))); let config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + FileScanConfigBuilder::new(StorageUrl::local_filesystem(), file_source) .with_file_groups(file_groups) .build(); @@ -3645,7 +3673,7 @@ mod tests { let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0))); let config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + FileScanConfigBuilder::new(StorageUrl::local_filesystem(), file_source) .with_file_groups(file_groups) .build(); @@ -3692,7 +3720,7 @@ mod tests { ]; let config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + FileScanConfigBuilder::new(StorageUrl::local_filesystem(), file_source) .with_file_groups(file_groups) .with_output_ordering(vec![ LexOrdering::new(vec![sort_expr.clone()]).unwrap(), @@ -3743,7 +3771,7 @@ mod tests { ])]; let config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + FileScanConfigBuilder::new(StorageUrl::local_filesystem(), file_source) .with_file_groups(file_groups) .with_output_ordering(vec![ LexOrdering::new(vec![sort_expr.clone()]).unwrap(), @@ -3809,7 +3837,7 @@ mod tests { ])]; let config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + FileScanConfigBuilder::new(StorageUrl::local_filesystem(), file_source) .with_file_groups(file_groups) .with_output_ordering(vec![ LexOrdering::new(vec![sort_expr.clone()]).unwrap(), @@ -3841,7 +3869,7 @@ mod tests { ])]; let config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + FileScanConfigBuilder::new(StorageUrl::local_filesystem(), file_source) .with_file_groups(file_groups) .with_output_ordering(vec![ LexOrdering::new(vec![sort_expr.clone()]).unwrap(), diff --git a/datafusion/datasource/src/file_scan_config/proto.rs b/datafusion/datasource/src/file_scan_config/proto.rs index d2dbb5f475479..763f67cd95050 100644 --- a/datafusion/datasource/src/file_scan_config/proto.rs +++ b/datafusion/datasource/src/file_scan_config/proto.rs @@ -41,7 +41,6 @@ use arrow::datatypes::Schema; use datafusion_common::parsers::CompressionTypeVariant; use datafusion_common::utils::{usize_from_wire, usize_to_wire}; use datafusion_common::{DataFusionError, Result, internal_datafusion_err}; -use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_physical_expr::Partitioning; use datafusion_physical_expr::projection::{ProjectionExpr, ProjectionExprs}; use datafusion_physical_expr_common::sort_expr::{ @@ -52,6 +51,7 @@ use datafusion_proto_models::datafusion_common::{ CompressionTypeVariant as ProtoCompressionTypeVariant, Schema as ProtoSchema, }; use datafusion_proto_models::protobuf; +use datafusion_storage::StorageUrl; use crate::file::FileSource; use crate::file_compression_type::FileCompressionType; @@ -74,6 +74,7 @@ impl FileScanConfig { // deciding how it is serialized is a compile error, not a silent // round-trip gap. let Self { + storage: _, object_store_url, file_groups, constraints, @@ -238,8 +239,8 @@ impl FileScanConfig { .collect::>>()?; let decoded_object_store_url = match object_store_url.is_empty() { - false => ObjectStoreUrl::parse(object_store_url)?, - true => ObjectStoreUrl::local_filesystem(), + false => StorageUrl::parse(object_store_url)?, + true => StorageUrl::local_filesystem(), }; let mut decoded_output_ordering = vec![]; @@ -319,8 +320,13 @@ impl FileScanConfig { "FileScanConfig: batch_size must be greater than 0" ); } + let storage = ctx + .task_ctx() + .runtime_env() + .storage(&decoded_object_store_url)?; let mut config = FileScanConfigBuilder::new(decoded_object_store_url, decoded_file_source) + .with_storage(storage) .with_file_groups(decoded_file_groups) .with_constraints(decoded_constraints) .with_statistics(decoded_statistics) diff --git a/datafusion/datasource/src/file_sink_config.rs b/datafusion/datasource/src/file_sink_config.rs index 48dce9a0cdb3e..7705d2bd85b2c 100644 --- a/datafusion/datasource/src/file_sink_config.rs +++ b/datafusion/datasource/src/file_sink_config.rs @@ -25,12 +25,12 @@ use crate::write::demux::{DemuxedStreamReceiver, start_demuxer_task}; use arrow::datatypes::{DataType, SchemaRef}; use datafusion_common::Result; use datafusion_common_runtime::SpawnedTask; -use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_expr::dml::InsertOp; +use datafusion_storage::StorageUrl; use async_trait::async_trait; -use object_store::ObjectStore; +use datafusion_storage::StorageBinding; #[cfg(feature = "proto")] mod proto; @@ -112,7 +112,7 @@ pub trait FileSink: DataSink { context: &Arc, demux_task: SpawnedTask>, file_stream_rx: DemuxedStreamReceiver, - object_store: Arc, + object_store: Arc, ) -> Result; /// File sink implementation of the [`DataSink::write_all`] method. @@ -122,9 +122,10 @@ pub trait FileSink: DataSink { context: &Arc, ) -> Result { let config = self.config(); - let object_store = context - .runtime_env() - .object_store(&config.object_store_url)?; + let object_store = match &config.storage { + Some(storage) => Arc::clone(storage), + None => context.runtime_env().storage(&config.object_store_url)?, + }; let (demux_task, file_stream_rx) = start_demuxer_task(config, data, context); self.spawn_writer_tasks_and_join( context, @@ -140,10 +141,12 @@ pub trait FileSink: DataSink { /// writing to any given file format. #[derive(Debug, Clone)] pub struct FileSinkConfig { + /// Immutable output registration retained by the physical plan. + pub storage: Option>, /// The unresolved URL specified by the user pub original_url: String, - /// Object store URL, used to get an ObjectStore instance - pub object_store_url: ObjectStoreUrl, + /// Storage namespace used when no immutable binding has been supplied. + pub object_store_url: StorageUrl, /// A collection of files organized into groups. /// Each FileGroup contains one or more PartitionedFile objects. pub file_group: FileGroup, diff --git a/datafusion/datasource/src/file_sink_config/proto.rs b/datafusion/datasource/src/file_sink_config/proto.rs index ed4b5c48bd2af..870dc7cfd7af6 100644 --- a/datafusion/datasource/src/file_sink_config/proto.rs +++ b/datafusion/datasource/src/file_sink_config/proto.rs @@ -20,9 +20,9 @@ use std::sync::Arc; use datafusion_common::{DataFusionError, Result, internal_datafusion_err}; -use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_expr::dml::InsertOp; use datafusion_proto_models::protobuf; +use datafusion_storage::StorageUrl; use crate::ListingTableUrl; use crate::file_groups::FileGroup; @@ -140,8 +140,9 @@ impl TryFrom<&protobuf::FileSinkConfig> for FileSinkConfig { })?; Ok(Self { + storage: None, original_url: String::default(), - object_store_url: ObjectStoreUrl::parse(&conf.object_store_url)?, + object_store_url: StorageUrl::parse(&conf.object_store_url)?, file_group, table_paths, output_schema: Arc::new(output_schema.try_into()?), @@ -162,7 +163,7 @@ mod tests { fn valid_file_sink_config() -> protobuf::FileSinkConfig { protobuf::FileSinkConfig { - object_store_url: ObjectStoreUrl::local_filesystem().to_string(), + object_store_url: StorageUrl::local_filesystem().to_string(), output_schema: Some( (&Schema::empty()) .try_into() diff --git a/datafusion/datasource/src/file_stream/mod.rs b/datafusion/datasource/src/file_stream/mod.rs index 619d2cf2ef479..3380ea4964f48 100644 --- a/datafusion/datasource/src/file_stream/mod.rs +++ b/datafusion/datasource/src/file_stream/mod.rs @@ -147,10 +147,9 @@ pub enum OnError { Skip, } -/// Generic API for opening a file using an [`ObjectStore`] and resolving to a +/// Generic API for opening a file using an [`StorageBinding`](datafusion_storage::StorageBinding) and resolving to a /// stream of [`RecordBatch`] /// -/// [`ObjectStore`]: object_store::ObjectStore pub trait FileOpener: Unpin + Send + Sync { /// Asynchronously open the specified file and return a stream /// of [`RecordBatch`] @@ -184,8 +183,8 @@ mod tests { use datafusion_common::DataFusionError; use datafusion_common::config::ConfigOptions; use datafusion_common::error::Result; - use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; + use datafusion_storage::StorageUrl; use futures::{FutureExt as _, StreamExt as _}; use std::collections::{BTreeMap, VecDeque}; use std::sync::Arc; @@ -318,7 +317,7 @@ mod tests { let table_schema = TableSchema::from(file_schema); let config = FileScanConfigBuilder::new( - ObjectStoreUrl::parse("test:///").unwrap(), + StorageUrl::parse("test:///").unwrap(), Arc::new(MockSource::new(table_schema)), ) .with_file_group(file_group) @@ -355,7 +354,7 @@ mod tests { fn builder_test_config() -> FileScanConfig { let table_schema = TableSchema::from(Arc::new(Schema::empty())); FileScanConfigBuilder::new( - ObjectStoreUrl::parse("test:///").unwrap(), + StorageUrl::parse("test:///").unwrap(), Arc::new(MockSource::new(table_schema)), ) .with_file(PartitionedFile::new("mock_file", 10)) @@ -1641,7 +1640,7 @@ mod tests { ) }); FileScanConfigBuilder::new( - ObjectStoreUrl::parse("test:///").unwrap(), + StorageUrl::parse("test:///").unwrap(), Arc::new(MockSource::new(table_schema)), ) .with_file_groups(file_groups) diff --git a/datafusion/datasource/src/mod.rs b/datafusion/datasource/src/mod.rs index d0bf64d81c71f..9059622936e48 100644 --- a/datafusion/datasource/src/mod.rs +++ b/datafusion/datasource/src/mod.rs @@ -65,8 +65,8 @@ use datafusion_common::stats::{Precision, is_known_empty}; use datafusion_common::{ColumnStatistics, Result, TableReference}; use datafusion_common::{ScalarValue, Statistics}; use datafusion_physical_expr::LexOrdering; +use datafusion_storage::{FileInfo, path::Path}; use futures::Stream; -use object_store::{ObjectMeta, path::Path}; pub use statistics::compute_all_files_statistics; use std::any::Any; use std::pin::Pin; @@ -124,7 +124,7 @@ impl FileRange { /// This enables query optimizers to use partition column bounds for pruning and planning. pub struct PartitionedFile { /// Path for the file (e.g. URL, filesystem path, etc) - pub object_meta: ObjectMeta, + pub object_meta: FileInfo, /// Values of partition columns to be appended to each row. /// /// These MUST have the same count, order, and type than the [`table_partition_cols`]. @@ -181,7 +181,7 @@ impl PartitionedFile { pub fn new(path: impl Into, size: u64) -> Self { Self { arrow_schema: None, - object_meta: ObjectMeta { + object_meta: FileInfo { location: Path::from(path.into()), last_modified: chrono::Utc.timestamp_nanos(0), size, @@ -198,8 +198,9 @@ impl PartitionedFile { } } - /// Create a file from a known ObjectMeta without partition - pub fn new_from_meta(object_meta: ObjectMeta) -> Self { + /// Create a file from a known FileInfo without partition + pub fn new_from_meta(object_meta: impl Into) -> Self { + let object_meta = object_meta.into(); Self { arrow_schema: None, object_meta, @@ -217,7 +218,7 @@ impl PartitionedFile { pub fn new_with_range(path: String, size: u64, start: i64, end: i64) -> Self { Self { arrow_schema: None, - object_meta: ObjectMeta { + object_meta: FileInfo { location: Path::from(path), last_modified: chrono::Utc.timestamp_nanos(0), size, @@ -395,8 +396,8 @@ impl PartitionedFile { } } -impl From for PartitionedFile { - fn from(object_meta: ObjectMeta) -> Self { +impl From for PartitionedFile { + fn from(object_meta: FileInfo) -> Self { PartitionedFile { object_meta, arrow_schema: None, @@ -468,7 +469,7 @@ pub fn generate_test_files(num_files: usize, overlap_factor: f64) -> Vec for protobuf::PartitionedFile { let last_modified = file.object_meta.last_modified; let last_modified_ns = last_modified.timestamp_nanos_opt().ok_or_else(|| { DataFusionError::Plan(format!( - "Invalid timestamp on PartitionedFile::ObjectMeta: {last_modified}" + "Invalid timestamp on PartitionedFile::FileInfo: {last_modified}" )) })? as u64; Ok(protobuf::PartitionedFile { @@ -97,7 +97,7 @@ impl TryFrom<&protobuf::PartitionedFile> for PartitionedFile { type Error = DataFusionError; fn try_from(file: &protobuf::PartitionedFile) -> Result { - let mut pf = PartitionedFile::new_from_meta(ObjectMeta { + let mut pf = PartitionedFile::new_from_meta(FileInfo { location: Path::parse(file.path.as_str()).map_err(|e| { internal_datafusion_err!("Invalid object_store path: {e}") })?, @@ -169,7 +169,7 @@ mod tests { #[test] fn partitioned_file_roundtrip_preserves_all_fields() -> Result<()> { let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); - let pf = PartitionedFile::new_from_meta(ObjectMeta { + let pf = PartitionedFile::new_from_meta(FileInfo { location: Path::parse("foo/bar.parquet")?, last_modified: Utc.timestamp_nanos(1_000_000_000), size: 1234, @@ -209,7 +209,7 @@ mod tests { // contains percent escapes must survive without a second round of // encoding or decoding. let path_str = "foo/foo%2Fbar/baz%252Fqux"; - let pf = PartitionedFile::new_from_meta(ObjectMeta { + let pf = PartitionedFile::new_from_meta(FileInfo { location: Path::parse(path_str)?, last_modified: Utc.timestamp_nanos(1_000), size: 42, diff --git a/datafusion/datasource/src/source.rs b/datafusion/datasource/src/source.rs index 741010c595197..5ac4c2899b171 100644 --- a/datafusion/datasource/src/source.rs +++ b/datafusion/datasource/src/source.rs @@ -26,10 +26,9 @@ use datafusion_physical_expr::projection::ProjectionExprs; use datafusion_physical_plan::execution_plan::{ Boundedness, EmissionType, SchedulingType, }; +use datafusion_physical_plan::metrics::BaselineMetrics; use datafusion_physical_plan::metrics::SplitMetrics; -use datafusion_physical_plan::metrics::{ - BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet, -}; +use datafusion_physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet}; use datafusion_physical_plan::projection::ProjectionExec; use datafusion_physical_plan::stream::BatchSplitStream; use datafusion_physical_plan::{ diff --git a/datafusion/datasource/src/test_util.rs b/datafusion/datasource/src/test_util.rs index 8b787aa104381..31ec656d7b2b3 100644 --- a/datafusion/datasource/src/test_util.rs +++ b/datafusion/datasource/src/test_util.rs @@ -25,7 +25,7 @@ use arrow::datatypes::Schema; use datafusion_common::{Result, tree_node::TreeNodeRecursion}; use datafusion_physical_expr::{PhysicalExpr, expressions::Column}; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; -use object_store::ObjectStore; +use datafusion_storage::StorageBinding; /// Minimal [`crate::file::FileSource`] implementation for use in tests. #[derive(Clone)] @@ -77,9 +77,10 @@ impl MockSource { impl FileSource for MockSource { fn create_file_opener( &self, - _object_store: Arc, + _object_store: Arc, _base_config: &FileScanConfig, _partition: usize, + _access_context: datafusion_storage::FileAccessContext, ) -> Result> { self.file_opener.clone().ok_or_else(|| { datafusion_common::internal_datafusion_err!("MockSource missing FileOpener") @@ -145,25 +146,70 @@ pub(crate) fn col(name: &str, schema: &Schema) -> Result> /// "one chunk containing everything", giving the single-chunk fast path. pub(crate) const CHUNK_SIZES: &[usize] = &[1, 2, 3, 4, 5, 7, 8, 11, 13, 16, usize::MAX]; -/// Seed a fresh `InMemory` store with `data` and wrap it in a -/// [`ChunkedStore`] that splits every GET response into `chunk_size`-byte -/// pieces. -pub(crate) async fn make_chunked_store( +/// A reader with controllable chunk boundaries for stream alignment tests. +pub(crate) fn make_chunked_reader( data: &[u8], chunk_size: usize, -) -> (Arc, object_store::path::Path) { - use bytes::Bytes; - use object_store::ObjectStoreExt; - use object_store::PutPayload; - use object_store::chunked::ChunkedStore; - use object_store::memory::InMemory; - use object_store::path::Path; - - let inner = Arc::new(InMemory::new()); - let path = Path::from("test"); - inner - .put(&path, PutPayload::from(Bytes::copy_from_slice(data))) - .await - .unwrap(); - (Arc::new(ChunkedStore::new(inner, chunk_size)), path) +) -> Arc { + #[derive(Debug)] + struct Reader { + data: bytes::Bytes, + chunk_size: usize, + } + #[async_trait::async_trait] + impl datafusion_storage::FileReader for Reader { + async fn read_range( + &self, + range: datafusion_storage::ReadRange, + ) -> datafusion_storage::Result { + let range = match range { + datafusion_storage::ReadRange::Bounded(r) => r, + datafusion_storage::ReadRange::Suffix(n) => { + (self.data.len() as u64).saturating_sub(n)..self.data.len() as u64 + } + }; + Ok(self.data.slice(range.start as usize..range.end as usize)) + } + + async fn read_ranges( + &self, + ranges: Vec>, + ) -> datafusion_storage::Result> { + Ok(ranges + .into_iter() + .map(|r| self.data.slice(r.start as usize..r.end as usize)) + .collect()) + } + fn stream( + self: Arc, + range: Option, + ) -> futures::stream::BoxStream<'static, datafusion_storage::Result> + { + use futures::StreamExt; + let size = self.data.len() as u64; + let range = match range { + Some(datafusion_storage::ReadRange::Bounded(r)) => r, + Some(datafusion_storage::ReadRange::Suffix(n)) => { + size.saturating_sub(n)..size + } + None => 0..size, + }; + futures::stream::unfold( + (self, range.start as usize, range.end as usize), + |(reader, start, end)| async move { + if start == end { + return None; + } + let next = start.saturating_add(reader.chunk_size).min(end); + let bytes = reader.data.slice(start..next); + Some((Ok(bytes), (reader, next, end))) + }, + ) + .boxed() + } + } + Arc::new(Reader { + data: bytes::Bytes::copy_from_slice(data), + chunk_size, + }) } diff --git a/datafusion/datasource/src/url.rs b/datafusion/datasource/src/url.rs index cfb6608ca0a78..ab7c6a1090361 100644 --- a/datafusion/datasource/src/url.rs +++ b/datafusion/datasource/src/url.rs @@ -20,17 +20,17 @@ use std::sync::Arc; use datafusion_common::{DataFusionError, Result, TableReference}; use datafusion_execution::cache::cache_manager::CachedFileList; use datafusion_execution::cache::cache_manager::TableScopedPath; -use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_session::Session; +use datafusion_storage::StorageUrl; +use datafusion_storage::path::DELIMITER; +use datafusion_storage::path::Path; +use datafusion_storage::{FileAccessContext, FileInfo, StorageBinding}; use futures::stream::BoxStream; use futures::{StreamExt, TryStreamExt}; use glob::Pattern; use itertools::Itertools; use log::debug; -use object_store::path::DELIMITER; -use object_store::path::Path; -use object_store::{ObjectMeta, ObjectStore, ObjectStoreExt}; use url::Url; /// A parsed URL identifying files for a listing table, see [`ListingTableUrl::parse`] @@ -247,10 +247,10 @@ impl ListingTableUrl { pub async fn list_prefixed_files<'a>( &'a self, ctx: &'a dyn Session, - store: &'a dyn ObjectStore, + store: &'a StorageBinding, prefix: Option, file_extension: &'a str, - ) -> Result>> { + ) -> Result>> { let exec_options = &ctx.config_options().execution; let ignore_subdirectory = exec_options.listing_table_ignore_subdirectory; @@ -263,7 +263,7 @@ impl ListingTableUrl { self.prefix.clone() }; - let list: BoxStream<'a, Result> = if self.is_collection() { + let list: BoxStream<'a, Result> = if self.is_collection() { list_with_cache( ctx, store, @@ -273,13 +273,16 @@ impl ListingTableUrl { ) .await? } else { - match store.head(&full_prefix).await { + match store + .stat(&full_prefix, &FileAccessContext::new("listing")) + .await + { Ok(meta) => futures::stream::once(async { Ok(meta) }) - .map_err(|e| DataFusionError::ObjectStore(Box::new(e))) + .map_err(|e| DataFusionError::Storage(Box::new(e))) .boxed(), // If the head command fails, it is likely that object doesn't exist. // Retry as though it were a prefix (aka a collection) - Err(object_store::Error::NotFound { .. }) => { + Err(datafusion_storage::Error::NotFound(_)) => { list_with_cache( ctx, store, @@ -307,9 +310,9 @@ impl ListingTableUrl { pub async fn list_all_files<'a>( &'a self, ctx: &'a dyn Session, - store: &'a dyn ObjectStore, + store: &'a StorageBinding, file_extension: &'a str, - ) -> Result>> { + ) -> Result>> { self.list_prefixed_files(ctx, store, None, file_extension) .await } @@ -319,10 +322,10 @@ impl ListingTableUrl { self.as_ref() } - /// Return the [`ObjectStoreUrl`] for this [`ListingTableUrl`] - pub fn object_store(&self) -> ObjectStoreUrl { + /// Return the [`StorageUrl`] for this [`ListingTableUrl`] + pub fn storage_url(&self) -> StorageUrl { let url = &self.url[url::Position::BeforeScheme..url::Position::BeforePath]; - ObjectStoreUrl::parse(url).unwrap() + StorageUrl::parse(url).unwrap() } /// Returns true if the [`ListingTableUrl`] points to the folder @@ -378,11 +381,11 @@ impl ListingTableUrl { /// subsequent prefix queries can be served from cache. async fn list_with_cache<'b>( ctx: &'b dyn Session, - store: &'b dyn ObjectStore, + store: &'b StorageBinding, table_ref: Option<&TableReference>, table_base_path: &Path, prefix: Option<&Path>, -) -> Result>> { +) -> Result>> { // Build the full listing path (table_base + prefix) let full_prefix = match prefix { Some(p) => { @@ -395,14 +398,15 @@ async fn list_with_cache<'b>( match ctx.runtime_env().cache_manager.get_list_files_cache() { None => Ok(store - .list(Some(&full_prefix)) - .map(|res| res.map_err(|e| DataFusionError::ObjectStore(Box::new(e)))) + .list(&full_prefix, FileAccessContext::new("listing")) + .map(|res| res.map_err(|e| DataFusionError::Storage(Box::new(e)))) .boxed()), Some(cache) => { // Build the filter prefix (only Some if prefix was requested) let filter_prefix = prefix.is_some().then(|| full_prefix.clone()); let table_scoped_base_path = TableScopedPath { + storage_id: store.id(), table: table_ref.cloned(), path: table_base_path.clone(), }; @@ -415,11 +419,11 @@ async fn list_with_cache<'b>( // Cache miss - always list and cache the full table // This ensures we have complete data for future prefix queries let mut vec = store - .list(Some(table_base_path)) - .try_collect::>() + .list(table_base_path, FileAccessContext::new("listing")) + .try_collect::>() .await?; vec.shrink_to_fit(); // Right-size before caching - let cached: CachedFileList = vec.into(); + let cached = CachedFileList::new(vec); let result = cached.files_matching_prefix(&filter_prefix); cache.put(&table_scoped_base_path, cached); result @@ -510,7 +514,6 @@ fn split_glob_expression(path: &str) -> Option<(&str, &str)> { mod tests { use super::*; use async_trait::async_trait; - use bytes::Bytes; use datafusion_common::DFSchema; use datafusion_common::config::TableOptions; use datafusion_execution::TaskContext; @@ -524,13 +527,8 @@ mod tests { use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use datafusion_physical_plan::ExecutionPlan; use datafusion_session::{CatalogProviderList, EmptyCatalogProviderList}; - use object_store::{ - CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, - PutMultipartOptions, PutPayload, - }; use std::any::Any; use std::collections::HashMap; - use std::ops::Range; use tempfile::tempdir; #[test] @@ -747,10 +745,7 @@ mod tests { #[tokio::test] async fn test_list_files() -> Result<()> { - let store = MockObjectStore { - in_mem: object_store::memory::InMemory::new(), - forbidden_paths: vec!["forbidden/e.parquet".into()], - }; + let store = mock_storage(vec!["forbidden/e.parquet".into()]); // Create some files: create_file(&store, "a.parquet").await; @@ -797,13 +792,13 @@ mod tests { ); // Including forbidden.parquet generates an error. - let Err(DataFusionError::ObjectStore(err)) = + let Err(DataFusionError::Storage(err)) = list_all_files("/forbidden/e.parquet", &store, "parquet").await else { panic!("Expected ObjectStore error"); }; - let object_store::Error::PermissionDenied { .. } = &*err else { + let datafusion_storage::Error::Backend { .. } = &*err else { panic!("Expected PermissionDenied error"); }; @@ -848,10 +843,7 @@ mod tests { async fn test_cache_path_equivalence() -> Result<()> { use datafusion_execution::runtime_env::RuntimeEnvBuilder; - let store = MockObjectStore { - in_mem: object_store::memory::InMemory::new(), - forbidden_paths: vec![], - }; + let store = mock_storage(vec![]); // Create test files with partition-style paths create_file(&store, "/table/year=2023/data1.parquet").await; @@ -954,10 +946,7 @@ mod tests { async fn test_cache_serves_partition_from_full_listing() -> Result<()> { use datafusion_execution::runtime_env::RuntimeEnvBuilder; - let store = MockObjectStore { - in_mem: object_store::memory::InMemory::new(), - forbidden_paths: vec![], - }; + let store = mock_storage(vec![]); // Create test files create_file(&store, "/sales/region=US/q1.parquet").await; @@ -1025,11 +1014,18 @@ mod tests { } /// Creates a file with "hello world" content at the specified path - async fn create_file(object_store: &dyn ObjectStore, path: &str) { - object_store - .put(&Path::from(path), PutPayload::from_static(b"hello world")) + async fn create_file(object_store: &StorageBinding, path: &str) { + use tokio::io::AsyncWriteExt; + let mut writer = object_store + .writer( + &Path::from(path), + datafusion_storage::WriterOptions::default(), + FileAccessContext::default(), + ) .await - .expect("failed to create test file"); + .unwrap(); + writer.write_all(b"hello world").await.unwrap(); + writer.shutdown().await.unwrap(); } /// Runs "list_prefixed_files" with no prefix to list all files and returns their paths @@ -1037,7 +1033,7 @@ mod tests { /// Panic's on error async fn list_all_files( url: &str, - store: &dyn ObjectStore, + store: &StorageBinding, file_extension: &str, ) -> Result> { try_list_prefixed_files(url, store, None, file_extension).await @@ -1048,7 +1044,7 @@ mod tests { /// Panic's on error async fn list_prefixed_files( url: &str, - store: &dyn ObjectStore, + store: &StorageBinding, prefix: Option, file_extension: &str, ) -> Result> { @@ -1058,7 +1054,7 @@ mod tests { /// Runs "list_prefixed_files" and returns their paths async fn try_list_prefixed_files( url: &str, - store: &dyn ObjectStore, + store: &StorageBinding, prefix: Option, file_extension: &str, ) -> Result> { @@ -1075,89 +1071,64 @@ mod tests { Ok(files) } - #[derive(Debug)] - struct MockObjectStore { - in_mem: object_store::memory::InMemory, - forbidden_paths: Vec, - } - - impl std::fmt::Display for MockObjectStore { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - self.in_mem.fmt(f) + fn mock_storage(forbidden_paths: Vec) -> Arc { + #[derive(Debug)] + struct Backend { + inner: Arc, + forbidden_paths: Vec, } - } - - #[async_trait] - impl ObjectStore for MockObjectStore { - async fn put_opts( - &self, - location: &Path, - payload: PutPayload, - opts: object_store::PutOptions, - ) -> object_store::Result { - self.in_mem.put_opts(location, payload, opts).await - } - - async fn put_multipart_opts( - &self, - location: &Path, - opts: PutMultipartOptions, - ) -> object_store::Result> { - self.in_mem.put_multipart_opts(location, opts).await - } - - async fn get_opts( - &self, - location: &Path, - options: GetOptions, - ) -> object_store::Result { - if options.head && self.forbidden_paths.contains(location) { - Err(object_store::Error::PermissionDenied { - path: location.to_string(), - source: "forbidden".into(), - }) - } else { - self.in_mem.get_opts(location, options).await + #[async_trait] + impl datafusion_storage::Storage for Backend { + async fn open( + &self, + file: &Path, + context: FileAccessContext, + ) -> datafusion_storage::Result> + { + self.inner.storage().open(file, context).await } - } - async fn get_ranges( - &self, - location: &Path, - ranges: &[Range], - ) -> object_store::Result> { - self.in_mem.get_ranges(location, ranges).await - } - - fn delete_stream( - &self, - locations: BoxStream<'static, object_store::Result>, - ) -> BoxStream<'static, object_store::Result> { - self.in_mem.delete_stream(locations) - } - - fn list( - &self, - prefix: Option<&Path>, - ) -> BoxStream<'static, object_store::Result> { - self.in_mem.list(prefix) - } - - async fn list_with_delimiter( - &self, - prefix: Option<&Path>, - ) -> object_store::Result { - self.in_mem.list_with_delimiter(prefix).await - } + async fn stat( + &self, + file: &Path, + context: &FileAccessContext, + ) -> datafusion_storage::Result { + if self.forbidden_paths.contains(file) { + return Err(datafusion_storage::Error::Backend { + backend: "test", + source: Box::new(std::io::Error::new( + std::io::ErrorKind::PermissionDenied, + "forbidden", + )), + }); + } + self.inner.storage().stat(file, context).await + } + fn list( + &self, + prefix: &Path, + context: FileAccessContext, + ) -> BoxStream<'_, datafusion_storage::Result> { + self.inner.storage().list(prefix, context) + } - async fn copy_opts( - &self, - from: &Path, - to: &Path, - options: CopyOptions, - ) -> object_store::Result<()> { - self.in_mem.copy_opts(from, to, options).await + async fn writer( + &self, + file: &Path, + options: datafusion_storage::WriterOptions, + context: FileAccessContext, + ) -> datafusion_storage::Result { + self.inner.storage().writer(file, options, context).await + } } + let inner = test_utils::storage::object_store(Arc::new( + object_store::memory::InMemory::new(), + )); + let backend = Arc::new(Backend { + inner, + forbidden_paths, + }); + Arc::new(StorageBinding::new(StorageUrl::local_filesystem(), backend)) } struct MockSession { diff --git a/datafusion/datasource/src/write/demux.rs b/datafusion/datasource/src/write/demux.rs index 1b3098d309789..b8c3ad8ecb3a9 100644 --- a/datafusion/datasource/src/write/demux.rs +++ b/datafusion/datasource/src/write/demux.rs @@ -43,8 +43,8 @@ use datafusion_common_runtime::SpawnedTask; use chrono::NaiveDate; use datafusion_execution::TaskContext; +use datafusion_storage::path::Path; use futures::StreamExt; -use object_store::path::Path; use rand::distr::SampleString; use tokio::sync::mpsc::{self, Receiver, Sender, UnboundedReceiver, UnboundedSender}; diff --git a/datafusion/datasource/src/write/mod.rs b/datafusion/datasource/src/write/mod.rs index c8c85112f0396..3e312759e2ae7 100644 --- a/datafusion/datasource/src/write/mod.rs +++ b/datafusion/datasource/src/write/mod.rs @@ -28,16 +28,16 @@ use datafusion_common::error::Result; use arrow::array::RecordBatch; use arrow::datatypes::Schema; use bytes::Bytes; -use object_store::ObjectStore; -use object_store::buffered::BufWriter; -use object_store::path::Path; +use datafusion_storage::StorageBinding; +use datafusion_storage::path::Path; +use datafusion_storage::{FileAccessContext, WriterOptions}; use tokio::io::AsyncWrite; pub mod demux; pub mod orchestration; /// A buffer with interior mutability shared by the SerializedFileWriter and -/// ObjectStore writer +/// storage writer #[derive(Clone)] pub struct SharedBuffer { /// The inner buffer for reading and writing @@ -101,30 +101,30 @@ pub fn get_writer_schema(config: &FileSinkConfig) -> Arc { /// /// This can be used to specify file compression on the writer. The writer /// will have a default buffer size unless altered. The specific default size -/// is chosen by [`BufWriter::new`]. +/// is chosen by the storage backend. /// -/// We drop the `AbortableWrite` struct and the writer will not try to cleanup on failure. -/// Users can configure automatic cleanup with their cloud provider. +/// Shutdown commits the output. Dropping it cancels pending I/O without committing; +/// remote cleanup follows the backend's policy. #[derive(Debug)] -pub struct ObjectWriterBuilder { +pub struct FileWriterBuilder { /// Compression type for object writer. file_compression_type: FileCompressionType, /// Output path location: Path, /// The related store that handles the given path - object_store: Arc, + object_store: Arc, /// The size of the buffer for the object writer. buffer_size: Option, /// The compression level for the object writer. compression_level: Option, } -impl ObjectWriterBuilder { - /// Create a new [`ObjectWriterBuilder`] for the specified path and compression type. +impl FileWriterBuilder { + /// Create a new [`FileWriterBuilder`] for the specified path and compression type. pub fn new( file_compression_type: FileCompressionType, location: &Path, - object_store: Arc, + object_store: Arc, ) -> Self { Self { file_compression_type, @@ -140,14 +140,14 @@ impl ObjectWriterBuilder { /// # Example /// ``` /// # use datafusion_datasource::file_compression_type::FileCompressionType; - /// # use datafusion_datasource::write::ObjectWriterBuilder; + /// # use datafusion_datasource::write::FileWriterBuilder; /// # use object_store::memory::InMemory; - /// # use object_store::path::Path; + /// # use datafusion_storage::path::Path; /// # use std::sync::Arc; /// # let compression_type = FileCompressionType::UNCOMPRESSED; /// # let location = Path::from("/foo/bar"); - /// # let object_store = Arc::new(InMemory::new()); - /// let mut builder = ObjectWriterBuilder::new(compression_type, &location, object_store); + /// # let object_store = test_utils::storage::object_store(Arc::new(InMemory::new())); + /// let mut builder = FileWriterBuilder::new(compression_type, &location, object_store); /// builder.set_buffer_size(Some(20 * 1024 * 1024)); //20 MiB /// assert_eq!( /// builder.get_buffer_size(), @@ -164,14 +164,14 @@ impl ObjectWriterBuilder { /// # Example /// ``` /// # use datafusion_datasource::file_compression_type::FileCompressionType; - /// # use datafusion_datasource::write::ObjectWriterBuilder; + /// # use datafusion_datasource::write::FileWriterBuilder; /// # use object_store::memory::InMemory; - /// # use object_store::path::Path; + /// # use datafusion_storage::path::Path; /// # use std::sync::Arc; /// # let compression_type = FileCompressionType::UNCOMPRESSED; /// # let location = Path::from("/foo/bar"); - /// # let object_store = Arc::new(InMemory::new()); - /// let builder = ObjectWriterBuilder::new(compression_type, &location, object_store) + /// # let object_store = test_utils::storage::object_store(Arc::new(InMemory::new())); + /// let builder = FileWriterBuilder::new(compression_type, &location, object_store) /// .with_buffer_size(Some(20 * 1024 * 1024)); //20 MiB /// assert_eq!( /// builder.get_buffer_size(), @@ -212,7 +212,7 @@ impl ObjectWriterBuilder { /// /// # Errors /// If there is an error applying the compression type. - pub fn build(self) -> Result> { + pub async fn build(self) -> Result> { let Self { file_compression_type, location, @@ -221,12 +221,14 @@ impl ObjectWriterBuilder { compression_level, } = self; - let buf_writer = match buffer_size { - Some(size) => BufWriter::with_capacity(object_store, location, size), - None => BufWriter::new(object_store, location), - }; + let writer = object_store + .writer( + &location, + WriterOptions { buffer_size }, + FileAccessContext::new("write"), + ) + .await?; - file_compression_type - .convert_async_writer_with_level(buf_writer, compression_level) + file_compression_type.convert_async_writer_with_level(writer, compression_level) } } diff --git a/datafusion/datasource/src/write/orchestration.rs b/datafusion/datasource/src/write/orchestration.rs index cd821b3b87897..967e36a7ab19a 100644 --- a/datafusion/datasource/src/write/orchestration.rs +++ b/datafusion/datasource/src/write/orchestration.rs @@ -22,7 +22,7 @@ use std::sync::Arc; use super::demux::DemuxedStreamReceiver; -use super::{BatchSerializer, ObjectWriterBuilder}; +use super::{BatchSerializer, FileWriterBuilder}; use crate::file_compression_type::FileCompressionType; use datafusion_common::error::Result; @@ -34,8 +34,8 @@ use datafusion_common_runtime::{JoinSet, SpawnedTask}; use datafusion_execution::TaskContext; use bytes::Bytes; +use datafusion_storage::StorageBinding; use futures::join; -use object_store::ObjectStore; use tokio::io::{AsyncWrite, AsyncWriteExt}; use tokio::sync::mpsc::{self, Receiver}; @@ -249,7 +249,7 @@ pub async fn spawn_writer_tasks_and_join( serializer: Arc, compression: FileCompressionType, compression_level: Option, - object_store: Arc, + object_store: Arc, demux_task: SpawnedTask>, mut file_stream_rx: DemuxedStreamReceiver, ) -> Result { @@ -266,7 +266,7 @@ pub async fn spawn_writer_tasks_and_join( }); while let Some((location, rb_stream)) = file_stream_rx.recv().await { let writer = - ObjectWriterBuilder::new(compression, &location, Arc::clone(&object_store)) + FileWriterBuilder::new(compression, &location, Arc::clone(&object_store)) .with_buffer_size(Some( context .session_config() @@ -275,7 +275,8 @@ pub async fn spawn_writer_tasks_and_join( .objectstore_writer_buffer_size, )) .with_compression_level(compression_level) - .build()?; + .build() + .await?; if tx_file_bundle .send((rb_stream, Arc::clone(&serializer), writer)) diff --git a/datafusion/execution/Cargo.toml b/datafusion/execution/Cargo.toml index c9d4acd3644ba..f802de47aa6ec 100644 --- a/datafusion/execution/Cargo.toml +++ b/datafusion/execution/Cargo.toml @@ -44,6 +44,7 @@ name = "datafusion_execution" default = ["sql"] parquet_encryption = [ + "parquet/encryption", ] arrow_buffer_pool = [ @@ -56,13 +57,14 @@ arrow = { workspace = true } arrow-buffer = { workspace = true } async-trait = { workspace = true } bytes = { workspace = true } +chrono = { workspace = true } dashmap = { workspace = true } datafusion-common = { workspace = true, default-features = false } datafusion-expr = { workspace = true, default-features = false } datafusion-physical-expr-common = { workspace = true, default-features = false } +datafusion-storage = { workspace = true } futures = { workspace = true } log = { workspace = true } -object_store = { workspace = true, features = ["fs"] } parking_lot = { workspace = true } parquet = { workspace = true, optional = true } pin-project-lite = { workspace = true } diff --git a/datafusion/execution/src/cache/cache_manager.rs b/datafusion/execution/src/cache/cache_manager.rs index bbbcaa4c77559..ea322ea0f45b7 100644 --- a/datafusion/execution/src/cache/cache_manager.rs +++ b/datafusion/execution/src/cache/cache_manager.rs @@ -21,8 +21,8 @@ use datafusion_common::HashMap; use datafusion_common::heap_size::{DFHeapSize, DFHeapSizeCtx}; use datafusion_common::{Result, Statistics}; use datafusion_physical_expr_common::sort_expr::LexOrdering; -use object_store::ObjectMeta; -use object_store::path::Path; +use datafusion_storage::FileInfo; +use datafusion_storage::path::Path; use std::any::Any; use std::fmt::{Debug, Formatter}; use std::ops::Deref; @@ -57,7 +57,7 @@ pub const DEFAULT_METADATA_CACHE_LIMIT: usize = 50 * 1024 * 1024; // 50M /// See [`crate::runtime_env::RuntimeEnv`] for more details pub type FileStatisticsCache = dyn Cache; -/// A cache for storing the [`ObjectMeta`]s that result from listing a path. +/// A cache for storing the [`FileInfo`]s that result from listing a path. /// /// Listing a path means doing an object store "list" operation or `ls` /// command on the local filesystem. This operation can be expensive, @@ -74,7 +74,7 @@ pub type ListFilesCache = dyn Cache; /// A cache for storing file-embedded metadata. /// /// This cache stores per-file metadata in the form of [`CachedFileMetadataEntry`], -/// which includes the [`ObjectMeta`] for validation. +/// which includes the [`FileInfo`] for validation. /// /// For example, the built in [`ListingTable`] uses this cache to avoid parsing /// Parquet footers multiple times for the same file. @@ -87,16 +87,17 @@ pub type ListFilesCache = dyn Cache; /// See [`crate::runtime_env::RuntimeEnv`] for more details. /// /// [`ListingTable`]: https://docs.rs/datafusion/latest/datafusion/datasource/listing/struct.ListingTable.html -pub type FileMetadataCache = dyn Cache; +pub type FileMetadataCache = + dyn Cache; /// Cached metadata for a file, including statistics and ordering. /// -/// This struct embeds the [`ObjectMeta`] used for cache validation, +/// This struct embeds the [`FileInfo`] used for cache validation, /// the `file_schema` fingerprint, cached statistics, and ordering information. #[derive(Debug, Clone, PartialEq, Eq)] pub struct CachedFileMetadata { /// File metadata used for cache validation (size, last_modified). - pub meta: ObjectMeta, + pub meta: FileInfo, /// Fingerprint of the `file_schema` used to compute `statistics`. pub schema_fingerprint: Arc, /// Cached statistics for the file, if available. @@ -108,7 +109,7 @@ pub struct CachedFileMetadata { impl CachedFileMetadata { /// Create a new cached file metadata entry. pub fn new( - meta: ObjectMeta, + meta: FileInfo, schema_fingerprint: Arc, statistics: Arc, ordering: Option, @@ -126,7 +127,7 @@ impl CachedFileMetadata { /// Returns true if the file size, last modified time, and schema match. pub fn is_valid_for( &self, - current_meta: &ObjectMeta, + current_meta: &FileInfo, current_schema_fingerprint: &Arc, ) -> bool { self.meta.size == current_meta.size @@ -163,19 +164,19 @@ impl DFHeapSize for CachedFileMetadata { #[derive(Debug, Clone, PartialEq)] pub struct CachedFileList { /// The cached file list. - pub files: Arc>, + pub files: Arc>, } impl CachedFileList { /// Create a new cached file list. - pub fn new(files: Vec) -> Self { + pub fn new(files: Vec) -> Self { Self { files: Arc::new(files), } } /// Filter the files by prefix. - fn filter_by_prefix(&self, prefix: Option<&Path>) -> Vec { + fn filter_by_prefix(&self, prefix: Option<&Path>) -> Vec { match prefix { Some(prefix) => self .files @@ -190,8 +191,8 @@ impl CachedFileList { /// Returns files matching the given prefix. /// /// When prefix is `None`, returns a clone of the `Arc` (no data copy). - /// When filtering is needed, returns a new `Arc` with filtered results (clones each matching [`ObjectMeta`]). - pub fn files_matching_prefix(&self, prefix: &Option) -> Arc> { + /// When filtering is needed, returns a new `Arc` with filtered results (clones each matching [`FileInfo`]). + pub fn files_matching_prefix(&self, prefix: &Option) -> Arc> { match prefix { None => Arc::clone(&self.files), Some(p) => Arc::new(self.filter_by_prefix(Some(p))), @@ -201,7 +202,7 @@ impl CachedFileList { impl CacheValue for CachedFileList { fn size(&self) -> usize { - self.files.capacity() * size_of::() + self.files.capacity() * size_of::() + self .files .iter() @@ -211,8 +212,8 @@ impl CacheValue for CachedFileList { } } -/// Calculates the number of bytes an [`ObjectMeta`] occupies in the heap. -pub fn meta_heap_bytes(object_meta: &ObjectMeta) -> usize { +/// Calculates the number of bytes an [`FileInfo`] occupies in the heap. +pub fn meta_heap_bytes(object_meta: &FileInfo) -> usize { let mut size = object_meta.location.as_ref().len(); if let Some(e) = &object_meta.e_tag { @@ -226,14 +227,14 @@ pub fn meta_heap_bytes(object_meta: &ObjectMeta) -> usize { } impl Deref for CachedFileList { - type Target = Arc>; + type Target = Arc>; fn deref(&self) -> &Self::Target { &self.files } } -impl From> for CachedFileList { - fn from(files: Vec) -> Self { +impl From> for CachedFileList { + fn from(files: Vec) -> Self { Self::new(files) } } @@ -260,7 +261,7 @@ pub trait FileMetadata: Any + Send + Sync { #[derive(Clone)] pub struct CachedFileMetadataEntry { /// File metadata used for cache validation (size, last_modified). - pub meta: ObjectMeta, + pub meta: FileInfo, /// The cached file metadata. pub file_metadata: Arc, } @@ -273,7 +274,7 @@ impl CacheValue for CachedFileMetadataEntry { impl CachedFileMetadataEntry { /// Create a new cached file metadata entry. - pub fn new(meta: ObjectMeta, file_metadata: Arc) -> Self { + pub fn new(meta: FileInfo, file_metadata: Arc) -> Self { Self { meta, file_metadata, @@ -281,7 +282,7 @@ impl CachedFileMetadataEntry { } /// Check if this cached entry is still valid for the given metadata. - pub fn is_valid_for(&self, current_meta: &ObjectMeta) -> bool { + pub fn is_valid_for(&self, current_meta: &FileInfo) -> bool { self.meta.size == current_meta.size && self.meta.last_modified == current_meta.last_modified } @@ -381,7 +382,7 @@ impl CacheManager { .map_or(0, |c| c.cache_limit()) } - /// Get the cache for storing the result of listing [`ObjectMeta`]s under the same path. + /// Get the cache for storing the result of listing [`FileInfo`]s under the same path. pub fn get_list_files_cache(&self) -> Option> { self.list_files_cache.clone() } diff --git a/datafusion/execution/src/cache/default_cache.rs b/datafusion/execution/src/cache/default_cache.rs index bfe326f3a47e1..51b927ef094d3 100644 --- a/datafusion/execution/src/cache/default_cache.rs +++ b/datafusion/execution/src/cache/default_cache.rs @@ -324,8 +324,8 @@ mod tests { use datafusion_expr::ColumnarValue; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; - use object_store::ObjectMeta; - use object_store::path::Path; + use datafusion_storage::FileInfo; + use datafusion_storage::path::Path; use std::sync::Mutex; use std::thread; use std::time::Duration; @@ -354,8 +354,8 @@ mod tests { } } - fn create_test_object_meta(path: &str, size: usize) -> ObjectMeta { - ObjectMeta { + fn create_test_object_meta(path: &str, size: usize) -> FileInfo { + FileInfo { location: Path::from(path), last_modified: DateTime::parse_from_rfc3339("2025-07-29T12:12:12+00:00") .unwrap() @@ -431,8 +431,8 @@ mod tests { fn generate_test_metadata_with_size( path: &str, size: usize, - ) -> (ObjectMeta, Arc) { - let object_meta = ObjectMeta { + ) -> (FileInfo, Arc) { + let object_meta = FileInfo { location: Path::from(path), last_modified: chrono::Utc::now(), size: size as u64, @@ -804,8 +804,8 @@ mod tests { assert_eq!(cache.list_entries(), HashMap::from([])); } - fn create_test_meta(path: &str, size: u64) -> ObjectMeta { - ObjectMeta { + fn create_test_meta(path: &str, size: u64) -> FileInfo { + FileInfo { location: Path::from(path), last_modified: DateTime::parse_from_rfc3339("2022-09-27T22:36:00+02:00") .unwrap() @@ -828,6 +828,7 @@ mod tests { )]); let path = TableScopedPath { + storage_id: 0, path: meta.location.clone(), table: None, }; @@ -889,6 +890,7 @@ mod tests { assert_eq!(entries.len(), 1); let path_3 = TableScopedPath { + storage_id: 0, path: Path::from("test"), table: None, }; @@ -964,6 +966,7 @@ mod tests { ); let path = TableScopedPath { + storage_id: 0, path: meta.location.clone(), table: None, }; @@ -993,6 +996,7 @@ mod tests { fn test_cache_invalidation_on_file_modification() { let cache = DefaultCache::new(DEFAULT_FILE_STATISTICS_MEMORY_LIMIT); let path = TableScopedPath { + storage_id: 0, path: Path::from("test.parquet"), table: None, }; @@ -1035,6 +1039,7 @@ mod tests { fn test_ordering_cache_invalidation_on_file_modification() { let cache = DefaultCache::new(DEFAULT_FILE_STATISTICS_MEMORY_LIMIT); let path = TableScopedPath { + storage_id: 0, path: Path::from("test.parquet"), table: None, }; @@ -1042,7 +1047,7 @@ mod tests { let schema_fingerprint = Arc::new(SchemaFingerprint::from_schema(&schema)); // Cache with original metadata and ordering - let meta_v1 = ObjectMeta { + let meta_v1 = FileInfo { location: path.path.clone(), last_modified: DateTime::parse_from_rfc3339("2022-09-27T22:36:00+02:00") .unwrap() @@ -1066,7 +1071,7 @@ mod tests { assert!(cached.ordering.is_some()); // File modified (size changed) - let meta_v2 = ObjectMeta { + let meta_v2 = FileInfo { location: path.path.clone(), last_modified: DateTime::parse_from_rfc3339("2022-09-28T10:00:00+02:00") .unwrap() @@ -1115,6 +1120,7 @@ mod tests { ); let path_1 = TableScopedPath { + storage_id: 0, path: meta1.location.clone(), table: None, }; @@ -1129,6 +1135,7 @@ mod tests { ); let path_2 = TableScopedPath { + storage_id: 0, path: meta2.location.clone(), table: None, }; @@ -1180,11 +1187,13 @@ mod tests { // create a cache with a limit which fits exactly 2 entries let cache = DefaultCache::new(limit_for_2_entries); let path_1 = TableScopedPath { + storage_id: 0, path: meta_1.location.clone(), table: None, }; let path_2 = TableScopedPath { + storage_id: 0, path: meta_2.location.clone(), table: None, }; @@ -1201,6 +1210,7 @@ mod tests { assert_eq!(result_2.unwrap(), value_2); let path_3 = TableScopedPath { + storage_id: 0, path: meta_3.location.clone(), table: None, }; @@ -1249,6 +1259,7 @@ mod tests { let cache = DefaultCache::new(limit_less_than_the_entry); let path_1 = TableScopedPath { + storage_id: 0, path: meta.location.clone(), table: None, }; @@ -1276,7 +1287,7 @@ mod tests { fn create_cached_file_metadata_with_stats( file_name: &str, series_size: i32, - ) -> (ObjectMeta, CachedFileMetadata) { + ) -> (FileInfo, CachedFileMetadata) { let series: Vec = (0..=series_size).collect(); let values = Int32Array::from(series); let offsets = OffsetBuffer::new(ScalarBuffer::from(vec![0, series_size + 1])); @@ -1335,8 +1346,8 @@ mod tests { } } - /// Helper function to create a test ObjectMeta with a specific path and location string size - fn create_object_meta(path: &str, location_size: usize) -> ObjectMeta { + /// Helper function to create a test FileInfo with a specific path and location string size + fn create_object_meta(path: &str, location_size: usize) -> FileInfo { // Create a location string of the desired size by padding with zeros let location_str = if location_size > path.len() { format!("{}{}", path, "0".repeat(location_size - path.len())) @@ -1344,7 +1355,7 @@ mod tests { path.to_string() }; - ObjectMeta { + FileInfo { location: Path::from(location_str), last_modified: DateTime::parse_from_rfc3339("2022-09-27T22:36:00+02:00") .unwrap() @@ -1363,10 +1374,11 @@ mod tests { table: Option, ) -> (TableScopedPath, CachedFileList) { let key = TableScopedPath { + storage_id: 0, table, path: Path::from(path), }; - let metas: Vec = (0..count) + let metas: Vec = (0..count) .map(|i| create_object_meta(&format!("file{i}"), meta_size)) .collect(); let value = CachedFileList::new(metas); @@ -1379,6 +1391,7 @@ mod tests { let table_ref = Some(TableReference::from("table")); let path = Path::from("test_path"); let key = TableScopedPath { + storage_id: 0, table: table_ref.clone(), path, }; @@ -1780,8 +1793,8 @@ mod tests { #[test] fn test_meta_heap_bytes_calculation() { - // Test with minimal ObjectMeta (no e_tag, no version) - let meta1 = ObjectMeta { + // Test with minimal FileInfo (no e_tag, no version) + let meta1 = FileInfo { location: Path::from("test"), last_modified: chrono::Utc::now(), size: 100, @@ -1791,7 +1804,7 @@ mod tests { assert_eq!(meta_heap_bytes(&meta1), 4); // Just the location string "test" // Test with e_tag - let meta2 = ObjectMeta { + let meta2 = FileInfo { location: Path::from("test"), last_modified: chrono::Utc::now(), size: 100, @@ -1801,7 +1814,7 @@ mod tests { assert_eq!(meta_heap_bytes(&meta2), 4 + 7); // location (4) + e_tag (7) // Test with version - let meta3 = ObjectMeta { + let meta3 = FileInfo { location: Path::from("test"), last_modified: chrono::Utc::now(), size: 100, @@ -1811,7 +1824,7 @@ mod tests { assert_eq!(meta_heap_bytes(&meta3), 4 + 4); // location (4) + version (4) // Test with both e_tag and version - let meta4 = ObjectMeta { + let meta4 = FileInfo { location: Path::from("test"), last_modified: chrono::Utc::now(), size: 100, @@ -1865,9 +1878,9 @@ mod tests { // Prefix filtering tests using CachedFileList::filter_by_prefix - /// Helper function to create ObjectMeta with a specific location path - fn create_object_meta_with_path(location: &str) -> ObjectMeta { - ObjectMeta { + /// Helper function to create FileInfo with a specific location path + fn create_object_meta_with_path(location: &str) -> FileInfo { + FileInfo { location: Path::from(location), last_modified: DateTime::parse_from_rfc3339("2022-09-27T22:36:00+02:00") .unwrap() @@ -1894,6 +1907,7 @@ mod tests { // Cache the full table listing let table_ref = Some(TableReference::from("table")); let key = TableScopedPath { + storage_id: 0, table: table_ref, path: table_base, }; @@ -1938,6 +1952,7 @@ mod tests { let table_ref = Some(TableReference::from("table")); let key = TableScopedPath { + storage_id: 0, table: table_ref, path: table_base, }; @@ -1972,6 +1987,7 @@ mod tests { let table_ref = Some(TableReference::from("table")); let key = TableScopedPath { + storage_id: 0, table: table_ref, path: table_base, }; diff --git a/datafusion/execution/src/cache/mod.rs b/datafusion/execution/src/cache/mod.rs index f47a3f3ca49f3..b1edc12b87cd5 100644 --- a/datafusion/execution/src/cache/mod.rs +++ b/datafusion/execution/src/cache/mod.rs @@ -24,9 +24,10 @@ use datafusion_common::arrow::datatypes::{DataType, Schema}; use datafusion_common::heap_size::{DFHeapSize, DFHeapSizeCtx}; use datafusion_common::instant::Instant; use datafusion_common::{HashMap, TableReference}; -use object_store::path::Path; +use datafusion_storage::path::Path; use std::collections::hash_map::DefaultHasher; -use std::fmt::{Debug, Display, Formatter}; +use std::fmt::Display; +use std::fmt::{Debug, Formatter}; use std::hash::{Hash, Hasher}; use std::time::Duration; @@ -148,6 +149,8 @@ impl CacheKey for TableScopedPath { /// table-level cache invalidation. #[derive(PartialEq, Eq, Hash, Clone, Debug)] pub struct TableScopedPath { + /// Isolates entries from registrations that reuse the same namespace and path. + pub storage_id: u64, pub table: Option, pub path: Path, } @@ -274,3 +277,23 @@ mod schema_fingerprint_tests { assert_eq!(plain, schema_md, "schema metadata must be ignored"); } } + +/// A path scoped to an immutable storage registration. +#[derive(PartialEq, Eq, Hash, Clone, Debug)] +pub struct FileCacheKey { + pub storage_id: u64, + pub path: Path, +} +impl CacheKey for FileCacheKey { + fn size(&self) -> usize { + self.path.size() + size_of::() + } + fn table_ref(&self) -> Option<&TableReference> { + None + } +} +impl Display for FileCacheKey { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "{}:{}", self.storage_id, self.path) + } +} diff --git a/datafusion/execution/src/lib.rs b/datafusion/execution/src/lib.rs index 5af7064f1cb8b..5404b283cd110 100644 --- a/datafusion/execution/src/lib.rs +++ b/datafusion/execution/src/lib.rs @@ -32,7 +32,6 @@ pub mod cache; pub mod config; pub mod disk_manager; pub mod memory_pool; -pub mod object_store; #[cfg(feature = "parquet_encryption")] pub mod parquet_encryption; pub mod runtime_env; diff --git a/datafusion/execution/src/object_store.rs b/datafusion/execution/src/object_store.rs deleted file mode 100644 index 22ce1f0cf2bbf..0000000000000 --- a/datafusion/execution/src/object_store.rs +++ /dev/null @@ -1,334 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -//! ObjectStoreRegistry holds all the object stores at Runtime with a scheme for each store. -//! This allows the user to extend DataFusion with different storage systems such as S3 or HDFS -//! and query data inside these systems. - -use dashmap::DashMap; -use datafusion_common::{ - DataFusionError, Result, exec_err, internal_datafusion_err, not_impl_err, -}; -use object_store::ObjectStore; -#[cfg(not(target_arch = "wasm32"))] -use object_store::local::LocalFileSystem; -use std::sync::Arc; -use url::Url; - -/// A parsed URL identifying a particular [`ObjectStore`] instance -/// -/// For example: -/// * `file://` for local file system -/// * `s3://bucket` for AWS S3 bucket -/// * `oss://bucket` for Aliyun OSS bucket -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct ObjectStoreUrl { - url: Url, -} - -impl ObjectStoreUrl { - /// Parse an [`ObjectStoreUrl`] from a string - /// - /// # Example - /// ``` - /// # use url::Url; - /// # use datafusion_execution::object_store::ObjectStoreUrl; - /// let object_store_url = ObjectStoreUrl::parse("s3://bucket").unwrap(); - /// assert_eq!(object_store_url.as_str(), "s3://bucket/"); - /// // can also access the underlying `Url` - /// let url: &Url = object_store_url.as_ref(); - /// assert_eq!(url.scheme(), "s3"); - /// assert_eq!(url.host_str(), Some("bucket")); - /// assert_eq!(url.path(), "/"); - /// ``` - pub fn parse(s: impl AsRef) -> Result { - let mut parsed = - Url::parse(s.as_ref()).map_err(|e| DataFusionError::External(Box::new(e)))?; - - let remaining = &parsed[url::Position::BeforePath..]; - if !remaining.is_empty() && remaining != "/" { - return exec_err!( - "ObjectStoreUrl must only contain scheme and authority, got: {remaining}" - ); - } - - // Always set path for consistency - parsed.set_path("/"); - Ok(Self { url: parsed }) - } - - /// An [`ObjectStoreUrl`] for the local filesystem (`file://`) - /// - /// # Example - /// ``` - /// # use datafusion_execution::object_store::ObjectStoreUrl; - /// let local_fs = ObjectStoreUrl::parse("file://").unwrap(); - /// assert_eq!(local_fs, ObjectStoreUrl::local_filesystem()) - /// ``` - pub fn local_filesystem() -> Self { - Self::parse("file://").unwrap() - } - - /// Returns this [`ObjectStoreUrl`] as a string - pub fn as_str(&self) -> &str { - self.as_ref() - } -} - -impl AsRef for ObjectStoreUrl { - fn as_ref(&self) -> &str { - self.url.as_ref() - } -} - -impl AsRef for ObjectStoreUrl { - fn as_ref(&self) -> &Url { - &self.url - } -} - -impl std::fmt::Display for ObjectStoreUrl { - fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { - self.as_str().fmt(f) - } -} - -/// [`ObjectStoreRegistry`] maps a URL to an [`ObjectStore`] instance, -/// and allows DataFusion to read from different [`ObjectStore`] -/// instances. For example DataFusion might be configured so that -/// -/// 1. `s3://my_bucket/lineitem/` mapped to the `/lineitem` path on an -/// AWS S3 object store bound to `my_bucket` -/// -/// 2. `s3://my_other_bucket/lineitem/` mapped to the (same) -/// `/lineitem` path on a *different* AWS S3 object store bound to -/// `my_other_bucket` -/// -/// When given a [`ListingTableUrl`], DataFusion tries to find an -/// appropriate [`ObjectStore`]. For example -/// -/// ```sql -/// create external table unicorns stored as parquet location 's3://my_bucket/lineitem/'; -/// ``` -/// -/// In this particular case, the url `s3://my_bucket/lineitem/` will be provided to -/// [`ObjectStoreRegistry::get_store`] and one of three things will happen: -/// -/// - If an [`ObjectStore`] has been registered with [`ObjectStoreRegistry::register_store`] with -/// `s3://my_bucket`, that [`ObjectStore`] will be returned -/// -/// - If an AWS S3 object store can be ad-hoc discovered by the url `s3://my_bucket/lineitem/`, this -/// object store will be registered with key `s3://my_bucket` and returned. -/// -/// - Otherwise an error will be returned, indicating that no suitable [`ObjectStore`] could -/// be found -/// -/// This allows for two different use-cases: -/// -/// 1. Systems where object store buckets are explicitly created using DDL, can register these -/// buckets using [`ObjectStoreRegistry::register_store`] -/// -/// 2. Systems relying on ad-hoc discovery, without corresponding DDL, can create [`ObjectStore`] -/// lazily by providing a custom implementation of [`ObjectStoreRegistry`] -/// -/// -/// [`ListingTableUrl`]: https://docs.rs/datafusion/latest/datafusion/datasource/listing/struct.ListingTableUrl.html -/// [`ObjectStore`]: object_store::ObjectStore -pub trait ObjectStoreRegistry: Send + Sync + std::fmt::Debug + 'static { - /// If a store with the same key existed before, it is replaced and returned - fn register_store( - &self, - url: &Url, - store: Arc, - ) -> Option>; - - /// Deregister the store previously registered with the same key. Returns the - /// deregistered store if it existed. - #[expect(unused_variables)] - fn deregister_store(&self, url: &Url) -> Result> { - not_impl_err!( - "ObjectStoreRegistry::deregister_store is not implemented for this ObjectStoreRegistry" - ) - } - - /// Get a suitable store for the provided URL. For example: - /// - /// - URL with scheme `file:///` or no scheme will return the default LocalFS store - /// - URL with scheme `s3://bucket/` will return the S3 store - /// - URL with scheme `hdfs://hostname:port/` will return the hdfs store - /// - /// If no [`ObjectStore`] found for the `url`, ad-hoc discovery may be executed depending on - /// the `url` and [`ObjectStoreRegistry`] implementation. An [`ObjectStore`] may be lazily - /// created and registered. - fn get_store(&self, url: &Url) -> Result>; -} - -/// The default [`ObjectStoreRegistry`] -pub struct DefaultObjectStoreRegistry { - /// A map from scheme to object store that serve list / read operations for the store - object_stores: DashMap>, -} - -impl std::fmt::Debug for DefaultObjectStoreRegistry { - fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { - f.debug_struct("DefaultObjectStoreRegistry") - .field( - "schemes", - &self - .object_stores - .iter() - .map(|o| o.key().clone()) - .collect::>(), - ) - .finish() - } -} - -impl Default for DefaultObjectStoreRegistry { - fn default() -> Self { - Self::new() - } -} - -impl DefaultObjectStoreRegistry { - /// This will register [`LocalFileSystem`] to handle `file://` paths - #[cfg(not(target_arch = "wasm32"))] - pub fn new() -> Self { - let object_stores: DashMap> = DashMap::new(); - object_stores.insert("file://".to_string(), Arc::new(LocalFileSystem::new())); - Self { object_stores } - } - - /// Default without any backend registered. - #[cfg(target_arch = "wasm32")] - pub fn new() -> Self { - let object_stores: DashMap> = DashMap::new(); - Self { object_stores } - } -} - -/// -/// Stores are registered based on the scheme, host and port of the provided URL -/// with a [`LocalFileSystem::new`] automatically registered for `file://` (if the -/// target arch is not `wasm32`). -/// -/// For example: -/// -/// - `file:///my_path` will return the default LocalFS store -/// - `s3://bucket/path` will return a store registered with `s3://bucket` if any -/// - `hdfs://host:port/path` will return a store registered with `hdfs://host:port` if any -impl ObjectStoreRegistry for DefaultObjectStoreRegistry { - fn register_store( - &self, - url: &Url, - store: Arc, - ) -> Option> { - let s = get_url_key(url); - self.object_stores.insert(s, store) - } - - fn deregister_store(&self, url: &Url) -> Result> { - let s = get_url_key(url); - let (_, object_store) = self.object_stores - .remove(&s) - .ok_or_else(|| { - internal_datafusion_err!("Failed to deregister object store. No suitable object store found for {url}. See `RuntimeEnv::register_object_store`") - })?; - - Ok(object_store) - } - - fn get_store(&self, url: &Url) -> Result> { - let s = get_url_key(url); - self.object_stores - .get(&s) - .map(|o| Arc::clone(o.value())) - .ok_or_else(|| { - internal_datafusion_err!("No suitable object store found for {url}. See `RuntimeEnv::register_object_store`") - }) - } -} - -/// Get the key of a url for object store registration. -/// The credential info will be removed -fn get_url_key(url: &Url) -> String { - format!( - "{}://{}", - url.scheme(), - &url[url::Position::BeforeHost..url::Position::AfterPort], - ) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_object_store_url() { - let file = ObjectStoreUrl::parse("file://").unwrap(); - assert_eq!(file.as_str(), "file:///"); - - let url = ObjectStoreUrl::parse("s3://bucket").unwrap(); - assert_eq!(url.as_str(), "s3://bucket/"); - - let url = ObjectStoreUrl::parse("s3://username:password@host:123").unwrap(); - assert_eq!(url.as_str(), "s3://username:password@host:123/"); - - let err = ObjectStoreUrl::parse("s3://bucket:invalid").unwrap_err(); - assert_eq!(err.strip_backtrace(), "External error: invalid port number"); - - let err = ObjectStoreUrl::parse("s3://bucket?").unwrap_err(); - assert_eq!( - err.strip_backtrace(), - "Execution error: ObjectStoreUrl must only contain scheme and authority, got: ?" - ); - - let err = ObjectStoreUrl::parse("s3://bucket?foo=bar").unwrap_err(); - assert_eq!( - err.strip_backtrace(), - "Execution error: ObjectStoreUrl must only contain scheme and authority, got: ?foo=bar" - ); - - let err = ObjectStoreUrl::parse("s3://host:123/foo").unwrap_err(); - assert_eq!( - err.strip_backtrace(), - "Execution error: ObjectStoreUrl must only contain scheme and authority, got: /foo" - ); - - let err = - ObjectStoreUrl::parse("s3://username:password@host:123/foo").unwrap_err(); - assert_eq!( - err.strip_backtrace(), - "Execution error: ObjectStoreUrl must only contain scheme and authority, got: /foo" - ); - } - - #[test] - fn test_get_url_key() { - let file = ObjectStoreUrl::parse("file://").unwrap(); - let key = get_url_key(&file.url); - assert_eq!(key.as_str(), "file://"); - - let url = ObjectStoreUrl::parse("s3://bucket").unwrap(); - let key = get_url_key(&url.url); - assert_eq!(key.as_str(), "s3://bucket"); - - let url = ObjectStoreUrl::parse("s3://username:password@host:123").unwrap(); - let key = get_url_key(&url.url); - assert_eq!(key.as_str(), "s3://host:123"); - } -} diff --git a/datafusion/execution/src/parquet_encryption.rs b/datafusion/execution/src/parquet_encryption.rs index 45eac10264e88..1b5abdb276d35 100644 --- a/datafusion/execution/src/parquet_encryption.rs +++ b/datafusion/execution/src/parquet_encryption.rs @@ -21,7 +21,7 @@ use dashmap::DashMap; use datafusion_common::config::EncryptionFactoryOptions; use datafusion_common::error::Result; use datafusion_common::internal_datafusion_err; -use object_store::path::Path; +use datafusion_storage::path::Path; use parquet::encryption::decrypt::FileDecryptionProperties; use parquet::encryption::encrypt::FileEncryptionProperties; use std::sync::Arc; diff --git a/datafusion/execution/src/runtime_env.rs b/datafusion/execution/src/runtime_env.rs index fcfe51267e65f..3c50ecd776c45 100644 --- a/datafusion/execution/src/runtime_env.rs +++ b/datafusion/execution/src/runtime_env.rs @@ -24,20 +24,19 @@ use crate::{ memory_pool::{ GreedyMemoryPool, MemoryPool, TrackConsumersPool, UnboundedMemoryPool, }, - object_store::{DefaultObjectStoreRegistry, ObjectStoreRegistry}, }; use crate::cache::cache_manager::{CacheManager, CacheManagerConfig}; #[cfg(feature = "parquet_encryption")] use crate::parquet_encryption::{EncryptionFactory, EncryptionFactoryRegistry}; use datafusion_common::{Result, config::ConfigEntry}; -use object_store::ObjectStore; +use std::path::PathBuf; use std::sync::Arc; +use std::time::Duration; use std::{ fmt::{Debug, Formatter}, num::NonZeroUsize, }; -use std::{path::PathBuf, time::Duration}; use url::Url; #[derive(Clone)] @@ -50,7 +49,7 @@ use url::Url; /// * [`MemoryPool`]: Manage memory /// * [`DiskManager`]: Manage temporary files on local disk /// * [`CacheManager`]: Manage temporary cache data during the session lifetime -/// * [`ObjectStoreRegistry`]: Manage mapping URLs to object store instances +/// * [`StorageRegistry`](datafusion_storage::StorageRegistry): Manage storage bindings /// /// # Example: Create default `RuntimeEnv` /// ``` @@ -71,14 +70,14 @@ use url::Url; /// .unwrap(); /// ``` pub struct RuntimeEnv { + /// Complete storage bindings shared by sessions using this runtime. + pub storage_registry: Arc, /// Runtime memory management pub memory_pool: Arc, /// Manage temporary files during query execution pub disk_manager: Arc, /// Manage temporary cache during query execution pub cache_manager: Arc, - /// Object Store Registry - pub object_store_registry: Arc, /// Parquet encryption factory registry #[cfg(feature = "parquet_encryption")] pub parquet_encryption_factory_registry: Arc, @@ -164,64 +163,31 @@ impl RuntimeConfigValues { } impl RuntimeEnv { - /// Registers a custom `ObjectStore` to be used with a specific url. - /// This allows DataFusion to create external tables from urls that do not have - /// built in support such as `hdfs://namenode:port/...`. - /// - /// Returns the [`ObjectStore`] previously registered for this - /// scheme, if any. - /// - /// See [`ObjectStoreRegistry`] for more details - /// - /// # Example: Register local file system object store - /// ``` - /// # use std::sync::Arc; - /// # use url::Url; - /// # use datafusion_execution::runtime_env::RuntimeEnv; - /// # let runtime_env = RuntimeEnv::default(); - /// let url = Url::try_from("file://").unwrap(); - /// let object_store = object_store::local::LocalFileSystem::new(); - /// // register the object store with the runtime environment - /// runtime_env.register_object_store(&url, Arc::new(object_store)); - /// ``` - /// - /// # Example: Register remote URL object store like [Github](https://github.com) - /// ``` - /// # use std::sync::Arc; - /// # use url::Url; - /// # use datafusion_execution::runtime_env::RuntimeEnv; - /// # let runtime_env = RuntimeEnv::default(); - /// # // use local store for example as http feature is not enabled - /// # let http_store = object_store::local::LocalFileSystem::new(); - /// // create a new object store via object_store::http::HttpBuilder; - /// let base_url = Url::parse("https://github.com").unwrap(); - /// // (note this example can't depend on the http feature) - /// // let http_store = HttpBuilder::new() - /// // .with_url(base_url.clone()) - /// // .build() - /// // .unwrap(); - /// // register the object store with the runtime environment - /// runtime_env.register_object_store(&base_url, Arc::new(http_store)); - /// ``` - pub fn register_object_store( + /// Atomically register or replace all capabilities of a storage backend. + pub fn register_storage( &self, url: &Url, - object_store: Arc, - ) -> Option> { - self.object_store_registry.register_store(url, object_store) + storage: Arc, + ) -> Result>> { + self.storage_registry + .register(url, storage) + .map_err(Into::into) } - /// Deregisters a custom `ObjectStore` previously registered for a specific url. - /// See [`ObjectStoreRegistry::deregister_store`] for more details. - pub fn deregister_object_store(&self, url: &Url) -> Result> { - self.object_store_registry.deregister_store(url) + /// Remove a binding. Plans already holding the binding remain valid. + pub fn deregister_storage( + &self, + url: &Url, + ) -> Result> { + self.storage_registry.deregister(url).map_err(Into::into) } - /// Retrieves a `ObjectStore` instance for a url by consulting the - /// registry. See [`ObjectStoreRegistry::get_store`] for more - /// details. - pub fn object_store(&self, url: impl AsRef) -> Result> { - self.object_store_registry.get_store(url.as_ref()) + /// Resolve the immutable storage binding for a namespace. + pub fn storage( + &self, + url: impl AsRef, + ) -> Result> { + Ok(self.storage_registry.get(url.as_ref())?) } /// Returns the current spilling progress @@ -304,42 +270,56 @@ impl RuntimeEnv { ) }; - let metadata_cache_limit = self.cache_manager.get_metadata_cache_limit(); - let metadata_cache_value = format_byte_size( - metadata_cache_limit - .try_into() - .expect("Metadata cache size conversion failed"), - ); - - let list_files_cache_limit = self.cache_manager.get_list_files_cache_limit(); - let list_files_cache_value = format_byte_size( - list_files_cache_limit - .try_into() - .expect("List files cache size conversion failed"), - ); - - let list_files_cache_ttl = self - .cache_manager - .get_list_files_cache_ttl() - .map(format_duration); - - let file_statistics_cache_limit = - self.cache_manager.get_file_statistic_cache_limit(); - let file_statistics_cache_value = format_byte_size( - file_statistics_cache_limit - .try_into() - .expect("File statistics cache size conversion failed"), - ); + let ( + metadata_cache_value, + list_files_cache_value, + list_files_cache_ttl, + file_statistics_cache_value, + ) = { + let metadata_cache_limit = self.cache_manager.get_metadata_cache_limit(); + let metadata_cache_value = format_byte_size( + metadata_cache_limit + .try_into() + .expect("Metadata cache size conversion failed"), + ); + + let list_files_cache_limit = self.cache_manager.get_list_files_cache_limit(); + let list_files_cache_value = format_byte_size( + list_files_cache_limit + .try_into() + .expect("List files cache size conversion failed"), + ); + + let list_files_cache_ttl = self + .cache_manager + .get_list_files_cache_ttl() + .map(format_duration); + + let file_statistics_cache_limit = + self.cache_manager.get_file_statistic_cache_limit(); + let file_statistics_cache_value = format_byte_size( + file_statistics_cache_limit + .try_into() + .expect("File statistics cache size conversion failed"), + ); + + ( + Some(metadata_cache_value), + Some(list_files_cache_value), + list_files_cache_ttl, + Some(file_statistics_cache_value), + ) + }; RuntimeConfigValues { memory_limit: memory_limit_value, max_temp_directory_size: Some(max_temp_dir_value), max_spill_merge_fan_in: Some(max_spill_merge_fan_in), temp_directory: temp_dir_value, - metadata_cache_limit: Some(metadata_cache_value), - list_files_cache_limit: Some(list_files_cache_value), + metadata_cache_limit: metadata_cache_value, + list_files_cache_limit: list_files_cache_value, list_files_cache_ttl, - file_statistics_cache_limit: Some(file_statistics_cache_value), + file_statistics_cache_limit: file_statistics_cache_value, } .into_config_entries() } @@ -356,6 +336,8 @@ impl Default for RuntimeEnv { /// See example on [`RuntimeEnv`] #[derive(Clone)] pub struct RuntimeEnvBuilder { + /// Complete storage bindings shared by sessions using this runtime. + pub storage_registry: Arc, /// DiskManager to manage temporary disk file usage pub disk_manager: Option>, /// DiskManager builder to manager temporary disk file usage @@ -366,8 +348,6 @@ pub struct RuntimeEnvBuilder { pub memory_pool: Option>, /// CacheManager to manage cache data pub cache_manager: CacheManagerConfig, - /// ObjectStoreRegistry to get object store based on url - pub object_store_registry: Arc, /// Parquet encryption factory registry #[cfg(feature = "parquet_encryption")] pub parquet_encryption_factory_registry: Arc, @@ -383,11 +363,11 @@ impl RuntimeEnvBuilder { /// New with default values pub fn new() -> Self { Self { + storage_registry: Arc::new(datafusion_storage::StorageRegistry::default()), disk_manager: Default::default(), disk_manager_builder: Default::default(), memory_pool: Default::default(), cache_manager: Default::default(), - object_store_registry: Arc::new(DefaultObjectStoreRegistry::default()), #[cfg(feature = "parquet_encryption")] parquet_encryption_factory_registry: Default::default(), } @@ -411,12 +391,12 @@ impl RuntimeEnvBuilder { self } - /// Customize object store registry - pub fn with_object_store_registry( + /// Share an existing storage registry across runtime environments. + pub fn with_storage_registry( mut self, - object_store_registry: Arc, + registry: Arc, ) -> Self { - self.object_store_registry = object_store_registry; + self.storage_registry = registry; self } @@ -483,11 +463,11 @@ impl RuntimeEnvBuilder { /// Build a RuntimeEnv pub fn build(self) -> Result { let Self { + storage_registry, disk_manager, disk_manager_builder, memory_pool, cache_manager, - object_store_registry, #[cfg(feature = "parquet_encryption")] parquet_encryption_factory_registry, } = self; @@ -501,10 +481,10 @@ impl RuntimeEnvBuilder { }; Ok(RuntimeEnv { + storage_registry, memory_pool, disk_manager, cache_manager: CacheManager::try_new(&cache_manager)?, - object_store_registry, #[cfg(feature = "parquet_encryption")] parquet_encryption_factory_registry, }) @@ -534,11 +514,11 @@ impl RuntimeEnvBuilder { }; Self { + storage_registry: Arc::clone(&runtime_env.storage_registry), disk_manager: Some(Arc::clone(&runtime_env.disk_manager)), disk_manager_builder: None, memory_pool: Some(Arc::clone(&runtime_env.memory_pool)), cache_manager: cache_config, - object_store_registry: Arc::clone(&runtime_env.object_store_registry), #[cfg(feature = "parquet_encryption")] parquet_encryption_factory_registry: Arc::clone( &runtime_env.parquet_encryption_factory_registry, diff --git a/datafusion/functions-aggregate/Cargo.toml b/datafusion/functions-aggregate/Cargo.toml index ae6bf70088f58..a55dc75ac4526 100644 --- a/datafusion/functions-aggregate/Cargo.toml +++ b/datafusion/functions-aggregate/Cargo.toml @@ -109,4 +109,5 @@ name = "variance" harness = false [features] +default = [] force_hash_collisions = ["datafusion-common/force_hash_collisions"] diff --git a/datafusion/functions-table/Cargo.toml b/datafusion/functions-table/Cargo.toml index fb02c2c5e2cb3..682f6f9ee24e7 100644 --- a/datafusion/functions-table/Cargo.toml +++ b/datafusion/functions-table/Cargo.toml @@ -40,6 +40,9 @@ workspace = true [lib] name = "datafusion_functions_table" +[features] +default = [] + [dependencies] arrow = { workspace = true } async-trait = { workspace = true } diff --git a/datafusion/functions/Cargo.toml b/datafusion/functions/Cargo.toml index 28765bdda9e2e..d0c1402d056ad 100644 --- a/datafusion/functions/Cargo.toml +++ b/datafusion/functions/Cargo.toml @@ -43,6 +43,7 @@ crypto_expressions = ["md-5", "sha2", "blake2", "blake3"] datetime_expressions = ["chrono-tz"] # Enable encoding by default so the doctests work. In general don't automatically enable all packages. default = [ + "datetime_expressions", "encoding_expressions", "math_expressions", @@ -111,8 +112,6 @@ required-features = ["math_expressions"] [[bench]] harness = false name = "misc" -required-features = [] - [[bench]] harness = false name = "regex_expressions" diff --git a/datafusion/physical-expr-adapter/Cargo.toml b/datafusion/physical-expr-adapter/Cargo.toml index 453c8bdaacb4a..0071eb0937b5e 100644 --- a/datafusion/physical-expr-adapter/Cargo.toml +++ b/datafusion/physical-expr-adapter/Cargo.toml @@ -15,6 +15,9 @@ rust-version = { workspace = true } name = "datafusion_physical_expr_adapter" path = "src/lib.rs" +[features] +default = [] + [dependencies] arrow = { workspace = true } datafusion-common = { workspace = true } diff --git a/datafusion/physical-optimizer/Cargo.toml b/datafusion/physical-optimizer/Cargo.toml index cb03303ac3c3f..fc35026c0f468 100644 --- a/datafusion/physical-optimizer/Cargo.toml +++ b/datafusion/physical-optimizer/Cargo.toml @@ -38,6 +38,7 @@ all-features = true workspace = true [features] +default = [] recursive_protection = ["dep:recursive"] [dependencies] diff --git a/datafusion/physical-plan/Cargo.toml b/datafusion/physical-plan/Cargo.toml index 534cea8ea9cbb..fce6e4119f4e6 100644 --- a/datafusion/physical-plan/Cargo.toml +++ b/datafusion/physical-plan/Cargo.toml @@ -38,6 +38,7 @@ all-features = true workspace = true [features] +default = [] force_hash_collisions = [] test_utils = ["arrow/test_utils"] tokio_coop = [] @@ -88,7 +89,7 @@ num-traits = { workspace = true } parking_lot = { workspace = true } pin-project-lite = { workspace = true } serde_json = { workspace = true, features = ["preserve_order"] } -tokio = { workspace = true } +tokio = { workspace = true, features = ["time"] } [dev-dependencies] arrow-data = { workspace = true } diff --git a/datafusion/proto/Cargo.toml b/datafusion/proto/Cargo.toml index 4b11bea01103f..e13b97f36ff69 100644 --- a/datafusion/proto/Cargo.toml +++ b/datafusion/proto/Cargo.toml @@ -74,7 +74,6 @@ datafusion-physical-expr-common = { workspace = true, features = ["proto"] } datafusion-physical-plan = { workspace = true, features = ["proto"] } datafusion-proto-common = { workspace = true } datafusion-proto-models = { workspace = true } -object_store = { workspace = true } prost = { workspace = true } recursive = { workspace = true, optional = true } serde_json = { workspace = true, optional = true } @@ -91,8 +90,11 @@ datafusion = { workspace = true, default-features = false, features = [ datafusion-functions = { workspace = true, default-features = true } datafusion-functions-aggregate = { workspace = true } datafusion-functions-window-common = { workspace = true } +datafusion-storage = { workspace = true } doc-comment = { workspace = true } flate2 = { workspace = true } +object_store = { workspace = true } pretty_assertions = "1.4" tempfile = { workspace = true } +test-utils = { path = "../../test-utils" } tokio = { workspace = true, features = ["rt-multi-thread"] } diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 887337e291f43..7b58a564391c8 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -119,7 +119,6 @@ mod file_scan_config_serde { }; use datafusion_datasource::file_stream::FileOpener; use datafusion_datasource::{PartitionedFile, TableSchema}; - use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_physical_expr::expressions::Column; use datafusion_physical_expr::projection::{ ProjectionExpr as FileProjectionExpr, ProjectionExprs as FileProjectionExprs, @@ -128,7 +127,7 @@ mod file_scan_config_serde { LexOrdering, Partitioning, PhysicalSortExpr, RangePartitioning, SplitPoint, }; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; - use object_store::ObjectStore; + use datafusion_storage::StorageUrl; #[derive(Clone)] struct SerdeTestSource { @@ -175,9 +174,10 @@ mod file_scan_config_serde { impl FileSource for SerdeTestSource { fn create_file_opener( &self, - _object_store: Arc, + _storage: Arc, _base_config: &FileScanConfig, _partition: usize, + _access_context: datafusion_storage::FileAccessContext, ) -> Result> { internal_err!("not needed for FileScanConfig serde tests") } @@ -275,7 +275,7 @@ mod file_scan_config_serde { ))]) .expect("single expression ordering"); - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source) + FileScanConfigBuilder::new(StorageUrl::local_filesystem(), source) .with_file_groups(vec![ FileGroup::new(vec![first_file, second_file]), FileGroup::new(vec![third_file]), @@ -322,10 +322,16 @@ mod file_scan_config_serde { impl FileScanSerdeHarness { fn new() -> Self { + let task_ctx = TaskContext::default(); + let binding = test_utils::storage::local(); + task_ctx + .runtime_env() + .register_storage(binding.url().as_ref(), Arc::clone(binding.storage())) + .unwrap(); Self { codec: DefaultPhysicalExtensionCodec {}, converter: DefaultPhysicalProtoConverter {}, - task_ctx: TaskContext::default(), + task_ctx, } } diff --git a/datafusion/proto/tests/cases/plans/dynamic_filters.rs b/datafusion/proto/tests/cases/plans/dynamic_filters.rs index 7892ffd9a1ab0..a93a7efb65cc1 100644 --- a/datafusion/proto/tests/cases/plans/dynamic_filters.rs +++ b/datafusion/proto/tests/cases/plans/dynamic_filters.rs @@ -24,7 +24,6 @@ use datafusion::arrow::compute::kernels::sort::SortOptions; use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use datafusion::datasource::empty::EmptyTable; use datafusion::datasource::listing::PartitionedFile; -use datafusion::datasource::object_store::ObjectStoreUrl; use datafusion::datasource::physical_plan::{ FileGroup, FileScanConfig, FileScanConfigBuilder, ParquetSource, }; @@ -51,6 +50,7 @@ use datafusion::physical_plan::{ PlanProperties, ReplaceChildrenOptions, SendableRecordBatchStream, }; use datafusion::prelude::SessionContext; +use datafusion::storage::StorageUrl; use datafusion_common::config::{ConfigOptions, TableParquetOptions}; use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{NullEquality, Result, internal_datafusion_err, internal_err}; @@ -171,7 +171,7 @@ fn roundtrip_dynamic_filter_plan_pair() -> Result<( .with_predicate(Arc::clone(&filter_expr_2)), ); let scan_config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + FileScanConfigBuilder::new(StorageUrl::local_filesystem(), file_source) .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( "/path/to/file.parquet".to_string(), 1024, @@ -381,10 +381,9 @@ fn datasource_for_dynamic_filter_pushdown( ParquetSource::new(Arc::clone(schema)) .with_table_parquet_options(parquet_options), ); - let scan_config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source) - .with_file(PartitionedFile::new("/path/to/file.parquet", 1024)) - .build(); + let scan_config = FileScanConfigBuilder::new(StorageUrl::local_filesystem(), source) + .with_file(PartitionedFile::new("/path/to/file.parquet", 1024)) + .build(); let mut config = ConfigOptions::default(); config.execution.parquet.pushdown_filters = true; diff --git a/datafusion/proto/tests/cases/plans/limits.rs b/datafusion/proto/tests/cases/plans/limits.rs index abcdd35fd840f..aeca3ebe551ca 100644 --- a/datafusion/proto/tests/cases/plans/limits.rs +++ b/datafusion/proto/tests/cases/plans/limits.rs @@ -22,7 +22,6 @@ use super::{roundtrip_test, roundtrip_test_and_return}; use datafusion::arrow::compute::kernels::sort::SortOptions; use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::datasource::listing::PartitionedFile; -use datafusion::datasource::object_store::ObjectStoreUrl; use datafusion::datasource::physical_plan::{ FileGroup, FileScanConfig, FileScanConfigBuilder, ParquetSource, }; @@ -42,6 +41,7 @@ use datafusion::physical_plan::{ ChildrenPropertiesMode, ExecutionPlan, ReplaceChildrenOptions, }; use datafusion::prelude::SessionContext; +use datafusion::storage::StorageUrl; use datafusion_common::Result; use datafusion_common::config::ConfigOptions; use datafusion_proto::physical_plan::{ @@ -186,7 +186,7 @@ fn roundtrip_limit_required_ordering_reaches_data_source() -> Result<()> { let make_scan = || { let file_source = Arc::new(ParquetSource::new(Arc::clone(&file_schema))); let scan_config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + FileScanConfigBuilder::new(StorageUrl::local_filesystem(), file_source) .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( "/path/to/file.parquet".to_string(), 1024, diff --git a/datafusion/proto/tests/cases/plans/sinks.rs b/datafusion/proto/tests/cases/plans/sinks.rs index 63bf11dd30f23..6e9732e24a81f 100644 --- a/datafusion/proto/tests/cases/plans/sinks.rs +++ b/datafusion/proto/tests/cases/plans/sinks.rs @@ -26,7 +26,6 @@ use datafusion::datasource::file_format::csv::CsvSink; use datafusion::datasource::file_format::json::JsonSink; use datafusion::datasource::file_format::parquet::ParquetSink; use datafusion::datasource::listing::{ListingTableUrl, PartitionedFile}; -use datafusion::datasource::object_store::ObjectStoreUrl; use datafusion::datasource::physical_plan::{ FileGroup, FileOutputMode, FileSink, FileSinkConfig, }; @@ -40,6 +39,7 @@ use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, SendableRecordBatchStream, }; use datafusion::prelude::SessionContext; +use datafusion::storage::StorageUrl; use datafusion_common::Result; use datafusion_common::config::TableParquetOptions; use datafusion_common::file_options::csv_writer::CsvWriterOptions; @@ -150,8 +150,9 @@ fn file_sink_config_roundtrip_preserves_fields() -> Result<()> { false, )])); let config = FileSinkConfig { + storage: None, original_url: "file:///tmp/output".to_string(), - object_store_url: ObjectStoreUrl::local_filesystem(), + object_store_url: StorageUrl::local_filesystem(), file_group: FileGroup::new(vec![PartitionedFile::new("/tmp/output", 1)]), table_paths: vec![ListingTableUrl::parse("file:///tmp/output")?], output_schema: schema, @@ -207,8 +208,9 @@ fn roundtrip_json_sink() -> Result<()> { let input = Arc::new(PlaceholderRowExec::new(schema.clone())); let file_sink_config = FileSinkConfig { + storage: None, original_url: String::default(), - object_store_url: ObjectStoreUrl::local_filesystem(), + object_store_url: StorageUrl::local_filesystem(), file_group: FileGroup::new(vec![PartitionedFile::new("/tmp".to_string(), 1)]), table_paths: vec![ListingTableUrl::parse("file:///")?], output_schema: schema.clone(), @@ -280,8 +282,9 @@ fn roundtrip_csv_sink() -> Result<()> { let input = Arc::new(PlaceholderRowExec::new(schema.clone())); let file_sink_config = FileSinkConfig { + storage: None, original_url: String::default(), - object_store_url: ObjectStoreUrl::local_filesystem(), + object_store_url: StorageUrl::local_filesystem(), file_group: FileGroup::new(vec![PartitionedFile::new("/tmp".to_string(), 1)]), table_paths: vec![ListingTableUrl::parse("file:///")?], output_schema: schema.clone(), @@ -333,8 +336,9 @@ fn roundtrip_parquet_sink() -> Result<()> { let input = Arc::new(PlaceholderRowExec::new(schema.clone())); let file_sink_config = FileSinkConfig { + storage: None, original_url: String::default(), - object_store_url: ObjectStoreUrl::local_filesystem(), + object_store_url: StorageUrl::local_filesystem(), file_group: FileGroup::new(vec![PartitionedFile::new("/tmp".to_string(), 1)]), table_paths: vec![ListingTableUrl::parse("file:///")?], output_schema: schema.clone(), diff --git a/datafusion/proto/tests/cases/plans/sources.rs b/datafusion/proto/tests/cases/plans/sources.rs index 7aff89eb3fdf7..38b7d37685ec6 100644 --- a/datafusion/proto/tests/cases/plans/sources.rs +++ b/datafusion/proto/tests/cases/plans/sources.rs @@ -31,7 +31,6 @@ use datafusion::datasource::file_format::json::JsonFormat; use datafusion::datasource::listing::{ ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl, PartitionedFile, }; -use datafusion::datasource::object_store::ObjectStoreUrl; use datafusion::datasource::physical_plan::{ ArrowSource, CsvSource, FileGroup, FileScanConfig, FileScanConfigBuilder, JsonSource, ParquetSource, wrap_partition_type_in_dict, wrap_partition_value_in_dict, @@ -50,6 +49,7 @@ use datafusion::physical_plan::{ }; use datafusion::prelude::SessionContext; use datafusion::scalar::ScalarValue; +use datafusion::storage::StorageUrl; use datafusion_common::config::TableParquetOptions; use datafusion_common::stats::Precision; use datafusion_common::{DataFusionError, Result, internal_datafusion_err, internal_err}; @@ -95,7 +95,7 @@ fn roundtrip_parquet_exec_with_pruning_predicate() -> Result<()> { ); let scan_config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + FileScanConfigBuilder::new(StorageUrl::local_filesystem(), file_source) .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( "/path/to/file.parquet".to_string(), 1024, @@ -142,7 +142,7 @@ async fn roundtrip_parquet_exec_with_sort_pushdown() -> Result<()> { fn file_scan_rejects_zero_batch_size() -> Result<()> { let schema = Arc::new(Schema::empty()); let scan_config = FileScanConfigBuilder::new( - ObjectStoreUrl::local_filesystem(), + StorageUrl::local_filesystem(), Arc::new(ParquetSource::new(schema)), ) .build(); @@ -174,12 +174,12 @@ fn file_scan_rejects_zero_batch_size() -> Result<()> { } #[test] -fn roundtrip_parquet_exec_attaches_cached_reader_factory_after_roundtrip() -> Result<()> { +fn roundtrip_parquet_exec_binds_storage_before_execution() -> Result<()> { let file_schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)])); let file_source = Arc::new(ParquetSource::new(Arc::clone(&file_schema))); let scan_config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + FileScanConfigBuilder::new(StorageUrl::local_filesystem(), file_source) .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( "/path/to/file.parquet".to_string(), 1024, @@ -216,10 +216,15 @@ fn roundtrip_parquet_exec_attaches_cached_reader_factory_after_roundtrip() -> Re internal_datafusion_err!("Expected ParquetSource after roundtrip") })?; - assert!( - parquet_source.parquet_file_reader_factory().is_some(), - "Parquet reader factory should be attached after decoding from protobuf" - ); + assert!(parquet_source.parquet_file_reader_factory().is_none()); + let binding = file_scan + .storage + .as_ref() + .expect("decoded scan has a storage binding"); + assert!(Arc::ptr_eq( + binding, + &ctx.runtime_env().storage(StorageUrl::local_filesystem())? + )); Ok(()) } @@ -248,7 +253,7 @@ fn roundtrip_arrow_scan() -> Result<()> { let file_source = Arc::new(ArrowSource::new_file_source(table_schema)); let scan_config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + FileScanConfigBuilder::new(StorageUrl::local_filesystem(), file_source) .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( "/path/to/file.arrow".to_string(), 1024, @@ -278,7 +283,7 @@ fn roundtrip_arrow_stream_scan() -> Result<()> { &file_schema, ))); let scan_config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + FileScanConfigBuilder::new(StorageUrl::local_filesystem(), file_source) .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( "/path/to/file.arrows".to_string(), 1024, @@ -305,7 +310,7 @@ fn arrow_scan_without_format_field_decodes_as_file_format() -> Result<()> { &file_schema, ))); let scan_config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + FileScanConfigBuilder::new(StorageUrl::local_filesystem(), file_source) .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( "/path/to/file.arrow".to_string(), 1024, @@ -349,7 +354,7 @@ fn roundtrip_json_scan_preserves_format_options() -> Result<()> { JsonSource::new(TableSchema::from(&file_schema)).with_newline_delimited(false), ); let scan_config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + FileScanConfigBuilder::new(StorageUrl::local_filesystem(), file_source) .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( "/path/to/file.json.gz".to_string(), 1024, @@ -466,7 +471,7 @@ fn roundtrip_avro_scan() -> Result<()> { Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)])); let file_source = Arc::new(AvroSource::new(TableSchema::from(&file_schema))); let scan_config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + FileScanConfigBuilder::new(StorageUrl::local_filesystem(), file_source) .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( "/path/to/file.avro".to_string(), 1024, @@ -496,7 +501,7 @@ fn roundtrip_csv_scan_preserves_format_options() -> Result<()> { })); let scan_config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + FileScanConfigBuilder::new(StorageUrl::local_filesystem(), file_source) .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( "/path/to/file.csv.gz".to_string(), 1024, @@ -573,7 +578,7 @@ async fn roundtrip_parquet_exec_with_table_partition_cols() -> Result<()> { let file_source = Arc::new(ParquetSource::new(table_schema.clone())); let scan_config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + FileScanConfigBuilder::new(StorageUrl::local_filesystem(), file_source) .with_projection_indices(Some(vec![0, 1]))? .with_file_group(FileGroup::new(vec![file_group])) .build(); @@ -596,7 +601,7 @@ fn roundtrip_parquet_exec_with_custom_predicate_expr() -> Result<()> { ); let scan_config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + FileScanConfigBuilder::new(StorageUrl::local_filesystem(), file_source) .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( "/path/to/file.parquet".to_string(), 1024, @@ -821,7 +826,7 @@ async fn roundtrip_projection_source() -> Result<()> { let file_source = Arc::new(ParquetSource::new(Arc::clone(&schema))); let scan_config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + FileScanConfigBuilder::new(StorageUrl::local_filesystem(), file_source) .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( "/path/to/file.parquet".to_string(), 1024, @@ -1025,7 +1030,7 @@ fn roundtrip_parquet_exec_output_partitioning() -> Result<()> { let output_partitioning = Partitioning::Hash(vec![Arc::new(Column::new("col", 0))], 1); let scan_config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + FileScanConfigBuilder::new(StorageUrl::local_filesystem(), file_source) .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( "/path/to/file.parquet".to_string(), 1024, @@ -1054,7 +1059,7 @@ fn roundtrip_parquet_exec_range_output_partitioning() -> Result<()> { vec![SplitPoint::new(vec![ScalarValue::Int32(Some(10))])], )); let scan_config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + FileScanConfigBuilder::new(StorageUrl::local_filesystem(), file_source) .with_file_groups(vec![ FileGroup::new(vec![PartitionedFile::new( "/path/to/file-1.parquet".to_string(), diff --git a/datafusion/pruning/Cargo.toml b/datafusion/pruning/Cargo.toml index a914a0a079bcd..1e54b9c71e30a 100644 --- a/datafusion/pruning/Cargo.toml +++ b/datafusion/pruning/Cargo.toml @@ -15,6 +15,9 @@ authors = { workspace = true } [lints] workspace = true +[features] +default = [] + [dependencies] arrow = { workspace = true } datafusion-common = { workspace = true, default-features = true } diff --git a/datafusion/session/Cargo.toml b/datafusion/session/Cargo.toml index e6c277a8b8493..5974ad0e0759e 100644 --- a/datafusion/session/Cargo.toml +++ b/datafusion/session/Cargo.toml @@ -30,6 +30,9 @@ version.workspace = true [package.metadata.docs.rs] all-features = true +[features] +default = [] + [dependencies] arrow-schema = { workspace = true } async-trait = { workspace = true } diff --git a/datafusion/sqllogictest/Cargo.toml b/datafusion/sqllogictest/Cargo.toml index c43f78f846092..70555dda920b9 100644 --- a/datafusion/sqllogictest/Cargo.toml +++ b/datafusion/sqllogictest/Cargo.toml @@ -40,6 +40,7 @@ workspace = true name = "datafusion_sqllogictest" [dependencies] +datafusion-storage.workspace = true arrow = { workspace = true } async-trait = { workspace = true } bigdecimal = { workspace = true } @@ -55,7 +56,6 @@ indicatif = "0.18" itertools = { workspace = true } log = { workspace = true } num-traits = { workspace = true } -object_store = { workspace = true } postgres-types = { version = "0.2.13", features = ["derive", "with-chrono-0_4"], optional = true } # When updating the following dependency verify that sqlite test file regeneration works correctly # by running the regenerate_sqlite_files.sh script. @@ -83,6 +83,7 @@ parquet_encryption = [ substrait = ["datafusion-substrait"] [dev-dependencies] +object_store = { workspace = true } env_logger = { workspace = true } regex = { workspace = true } # Required to make sure tests for pg display behaves consistently diff --git a/datafusion/sqllogictest/src/engines/datafusion_engine/normalize.rs b/datafusion/sqllogictest/src/engines/datafusion_engine/normalize.rs index 6172ffe35026b..b8f115500839d 100644 --- a/datafusion/sqllogictest/src/engines/datafusion_engine/normalize.rs +++ b/datafusion/sqllogictest/src/engines/datafusion_engine/normalize.rs @@ -146,7 +146,7 @@ fn normalize_paths(mut row: Vec) -> Vec { } /// The location of the datafusion checkout -static WORKSPACE_ROOT: LazyLock = LazyLock::new(|| { +static WORKSPACE_ROOT: LazyLock = LazyLock::new(|| { // e.g. /Software/datafusion/datafusion/core let dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")); @@ -162,12 +162,15 @@ static WORKSPACE_ROOT: LazyLock = LazyLock::new(|| { let sanitized_workplace_root = if cfg!(windows) { // Object store paths are delimited with `/`, e.g. `/datafusion/datafusion/testing/data/csv/aggregate_test_100.csv`. // The default windows delimiter is `\`, so the workplace path is `datafusion\datafusion`. - workspace_root.replace(std::path::MAIN_SEPARATOR, object_store::path::DELIMITER) + workspace_root.replace( + std::path::MAIN_SEPARATOR, + datafusion_storage::path::DELIMITER, + ) } else { workspace_root.to_string() }; - object_store::path::Path::parse(sanitized_workplace_root).unwrap() + datafusion_storage::path::Path::parse(sanitized_workplace_root).unwrap() }); macro_rules! get_row_value { diff --git a/datafusion/storage-object-store/Cargo.toml b/datafusion/storage-object-store/Cargo.toml new file mode 100644 index 0000000000000..36dde99e22f28 --- /dev/null +++ b/datafusion/storage-object-store/Cargo.toml @@ -0,0 +1,41 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +[package] +name = "datafusion-storage-object-store" +description = "object_store adapter for DataFusion storage" +version.workspace = true +edition.workspace = true +license.workspace = true +rust-version.workspace = true +repository.workspace = true + +[lints] +workspace = true + +[dependencies] +async-trait.workspace = true +bytes.workspace = true +chrono.workspace = true +datafusion-storage.workspace = true +futures.workspace = true +object_store = { workspace = true, features = ["fs"] } +tokio.workspace = true + +[dev-dependencies] +tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } diff --git a/datafusion/storage-object-store/src/lib.rs b/datafusion/storage-object-store/src/lib.rs new file mode 100644 index 0000000000000..693afb79148c8 --- /dev/null +++ b/datafusion/storage-object-store/src/lib.rs @@ -0,0 +1,207 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Adapter for the Arrow `object_store` crate. +use ::object_store::{GetOptions, GetRange, ObjectStore, ObjectStoreExt}; +use async_trait::async_trait; +use bytes::Bytes; +use datafusion_storage::path::Path; +use datafusion_storage::*; +use futures::{StreamExt, TryStreamExt, stream::BoxStream}; +use std::{ops::Range, sync::Arc}; + +/// Wrap an existing client; credentials, pools, and middleware stay on that client. +#[derive(Debug)] +pub struct ObjectStoreStorage(pub Arc); + +impl ObjectStoreStorage { + pub fn new(store: Arc) -> Self { + Self(store) + } + + #[cfg(not(target_arch = "wasm32"))] + pub fn local() -> Self { + Self(Arc::new(::object_store::local::LocalFileSystem::new())) + } +} + +fn error(error: ::object_store::Error) -> Error { + match error { + ::object_store::Error::NotFound { path, .. } => Error::NotFound(path), + other => Error::Backend { + backend: "object_store", + source: Box::new(other), + }, + } +} + +fn location(path: &Path) -> Result<::object_store::path::Path> { + ::object_store::path::Path::parse(path.as_ref()).map_err(|e| error(e.into())) +} + +fn get_range(range: ReadRange) -> GetRange { + match range { + ReadRange::Bounded(range) => GetRange::Bounded(range), + ReadRange::Suffix(size) => GetRange::Suffix(size), + } +} + +#[async_trait] +impl Storage for ObjectStoreStorage { + async fn open( + &self, + path: &Path, + _context: FileAccessContext, + ) -> Result> { + Ok(Arc::new(Reader { + store: self.0.clone(), + path: location(path)?, + })) + } + + async fn stat(&self, path: &Path, _context: &FileAccessContext) -> Result { + Ok(file_info( + self.0.head(&location(path)?).await.map_err(error)?, + )) + } + + async fn list_with_delimiter( + &self, + prefix: &Path, + _context: &FileAccessContext, + ) -> Result { + let result = self + .0 + .list_with_delimiter(Some(&location(prefix)?)) + .await + .map_err(error)?; + Ok(DirectoryListing { + files: result.objects.into_iter().map(file_info).collect(), + directories: result + .common_prefixes + .into_iter() + .map(|path| Path::parse(path.as_ref())) + .collect::>()?, + }) + } + + fn list( + &self, + prefix: &Path, + _context: FileAccessContext, + ) -> BoxStream<'_, Result> { + let path = match location(prefix) { + Ok(path) => path, + Err(e) => return futures::stream::once(async { Err(e) }).boxed(), + }; + self.0 + .list(Some(&path)) + .map(|meta| meta.map(file_info).map_err(error)) + .boxed() + } + + async fn writer( + &self, + path: &Path, + options: WriterOptions, + _context: FileAccessContext, + ) -> Result { + let path = location(path)?; + let writer = match options.buffer_size { + Some(size) => ::object_store::buffered::BufWriter::with_capacity( + self.0.clone(), + path, + size, + ), + None => ::object_store::buffered::BufWriter::new(self.0.clone(), path), + }; + Ok(Box::new(writer)) + } +} + +#[derive(Debug)] +struct Reader { + store: Arc, + path: ::object_store::path::Path, +} + +#[async_trait] +impl FileReader for Reader { + async fn read_range(&self, range: ReadRange) -> Result { + match range { + ReadRange::Bounded(range) => { + self.store.get_range(&self.path, range).await.map_err(error) + } + range @ ReadRange::Suffix(_) => self + .store + .get_opts( + &self.path, + GetOptions { + range: Some(get_range(range)), + ..Default::default() + }, + ) + .await + .map_err(error)? + .bytes() + .await + .map_err(error), + } + } + + async fn read_ranges(&self, ranges: Vec>) -> Result> { + self.store + .get_ranges(&self.path, &ranges) + .await + .map_err(error) + } + + fn stream( + self: Arc, + range: Option, + ) -> BoxStream<'static, Result> { + futures::stream::once(async move { + Ok::<_, Error>( + self.store + .get_opts( + &self.path, + GetOptions { + range: range.map(get_range), + ..Default::default() + }, + ) + .await + .map_err(error)? + .into_stream() + .map_err(error), + ) + }) + .try_flatten() + .boxed() + } +} + +/// Convert metadata at the SDK boundary, retaining both ETag and version. +pub fn file_info(meta: ::object_store::ObjectMeta) -> FileInfo { + FileInfo { + location: Path::parse(meta.location.as_ref()).expect("validated SDK path"), + size: meta.size, + last_modified: meta.last_modified, + e_tag: meta.e_tag, + version: meta.version, + } +} diff --git a/datafusion/storage-object-store/tests/contracts.rs b/datafusion/storage-object-store/tests/contracts.rs new file mode 100644 index 0000000000000..f61d9461e3375 --- /dev/null +++ b/datafusion/storage-object-store/tests/contracts.rs @@ -0,0 +1,73 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#[path = "../../storage/tests/common/backend.rs"] +mod backend; +use datafusion_storage::{path::Path, *}; +use datafusion_storage_object_store::ObjectStoreStorage; +use std::sync::Arc; + +#[tokio::test] +async fn object_store_adapter_preserves_unconditional_reads() { + use object_store::{ObjectStoreExt, memory::InMemory}; + let backend = Arc::new(InMemory::new()); + let storage = ObjectStoreStorage::new(backend.clone()); + backend::backend_contract(&storage).await; + let path = Path::from("data.parquet"); + let metadata = storage + .stat(&path, &FileAccessContext::default()) + .await + .unwrap(); + assert!(metadata.e_tag.is_some()); + let reader = storage + .open(&path, FileAccessContext::default()) + .await + .unwrap(); + backend + .put( + &object_store::path::Path::from("data.parquet"), + bytes::Bytes::from_static(b"replacement").into(), + ) + .await + .unwrap(); + assert_eq!(reader.read_range((0..3).into()).await.unwrap(), b"rep"[..]); + assert_eq!( + reader.read_ranges(vec![0..1, 2..3]).await.unwrap(), + vec![ + bytes::Bytes::from_static(b"r"), + bytes::Bytes::from_static(b"p") + ] + ); + assert_eq!( + collect_bytes(reader.stream(None)).await.unwrap(), + b"replacement"[..] + ); +} + +#[test] +fn file_metadata_keeps_both_version_and_etag() { + let metadata = object_store::ObjectMeta { + location: object_store::path::Path::from("file"), + last_modified: chrono::DateTime::UNIX_EPOCH, + size: 42, + e_tag: Some("etag".into()), + version: Some("version".into()), + }; + let info = datafusion_storage_object_store::file_info(metadata); + assert_eq!(info.e_tag.as_deref(), Some("etag")); + assert_eq!(info.version.as_deref(), Some("version")); +} diff --git a/datafusion/storage-opendal/Cargo.toml b/datafusion/storage-opendal/Cargo.toml new file mode 100644 index 0000000000000..1f5a1f0add535 --- /dev/null +++ b/datafusion/storage-opendal/Cargo.toml @@ -0,0 +1,42 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +[package] +name = "datafusion-storage-opendal" +description = "opendal adapter for DataFusion storage" +version.workspace = true +edition.workspace = true +license.workspace = true +rust-version.workspace = true +repository.workspace = true + +[lints] +workspace = true + +[dependencies] +async-trait.workspace = true +bytes.workspace = true +chrono.workspace = true +datafusion-storage.workspace = true +futures.workspace = true +opendal = { workspace = true } +tokio-util = { workspace = true, features = ["compat"] } + +[dev-dependencies] +opendal = { workspace = true, features = ["services-memory"] } +tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } diff --git a/datafusion/storage-opendal/src/lib.rs b/datafusion/storage-opendal/src/lib.rs new file mode 100644 index 0000000000000..2b1988fee5345 --- /dev/null +++ b/datafusion/storage-opendal/src/lib.rs @@ -0,0 +1,235 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Direct OpenDAL adapter. No ObjectStore bridge is involved. +use async_trait::async_trait; +use bytes::Bytes; +use chrono::DateTime; +use datafusion_storage::{path::Path, *}; +use futures::{StreamExt, TryStreamExt, stream::BoxStream}; +use std::{fmt::Debug, ops::Range, sync::Arc}; +use tokio_util::compat::FuturesAsyncWriteCompatExt; + +/// Wrap an existing operator, including its layers and credential configuration. +#[derive(Debug)] +pub struct OpendalStorage(pub ::opendal::Operator); + +impl OpendalStorage { + pub fn new(operator: ::opendal::Operator) -> Self { + Self(operator) + } +} + +fn error(error: ::opendal::Error) -> Error { + match error.kind() { + ::opendal::ErrorKind::NotFound => Error::NotFound(error.to_string()), + _ => Error::Backend { + backend: "opendal", + source: Box::new(error), + }, + } +} + +fn file_info(path: Path, meta: &::opendal::Metadata) -> FileInfo { + FileInfo { + location: path, + size: meta.content_length(), + last_modified: meta + .last_modified() + .and_then(|t| { + DateTime::from_timestamp( + t.into_inner().as_second(), + t.into_inner().subsec_nanosecond() as u32, + ) + }) + .unwrap_or(DateTime::UNIX_EPOCH), + e_tag: meta.etag().map(String::from), + version: meta.version().map(String::from), + } +} + +#[async_trait] +impl Storage for OpendalStorage { + async fn open( + &self, + path: &Path, + _context: FileAccessContext, + ) -> Result> { + Ok(Arc::new(Reader { + inner: self.0.reader(path.as_ref()).await.map_err(error)?, + operator: self.0.clone(), + path: path.clone(), + })) + } + + async fn stat(&self, path: &Path, _context: &FileAccessContext) -> Result { + Ok(file_info( + path.clone(), + &self.0.stat(path.as_ref()).await.map_err(error)?, + )) + } + + async fn list_with_delimiter( + &self, + prefix: &Path, + context: &FileAccessContext, + ) -> Result { + let mut entries = self + .0 + .lister(&directory_path(prefix)) + .await + .map_err(error)?; + let mut result = DirectoryListing::default(); + while let Some(entry) = entries.try_next().await.map_err(error)? { + let path = Path::parse(entry.path())?; + if entry.metadata().is_dir() { + if &path != prefix { + result.directories.push(path); + } + } else { + result.files.push(self.stat(&path, context).await?); + } + } + Ok(result) + } + + fn list( + &self, + prefix: &Path, + context: FileAccessContext, + ) -> BoxStream<'_, Result> { + let path = directory_path(prefix); + futures::stream::once(async move { + Ok::<_, Error>( + self.0 + .lister_with(&path) + .recursive(true) + .await + .map_err(error)? + .map_err(error), + ) + }) + .try_flatten() + .try_filter(|entry| futures::future::ready(entry.metadata().is_file())) + .and_then(move |entry| { + let context = context.clone(); + async move { self.stat(&Path::parse(entry.path())?, &context).await } + }) + .boxed() + } + + async fn writer( + &self, + path: &Path, + options: WriterOptions, + _context: FileAccessContext, + ) -> Result { + let mut builder = self.0.writer_with(path.as_ref()); + if let Some(size) = options.buffer_size { + builder = builder.chunk(size); + } + Ok(Box::new( + builder + .await + .map_err(error)? + .into_futures_async_write() + .compat_write(), + )) + } +} + +struct Reader { + inner: ::opendal::Reader, + operator: ::opendal::Operator, + path: Path, +} +impl Debug for Reader { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("OpendalReader") + .field("path", &self.path) + .finish_non_exhaustive() + } +} +impl Reader { + async fn range(&self, range: ReadRange) -> Result> { + match range { + ReadRange::Bounded(range) => Ok(range), + ReadRange::Suffix(length) => { + // OpenDAL reads offset ranges. Resolve size only for suffix requests. + let size = match self.inner.metadata() { + Some(meta) => meta.content_length(), + None => self + .operator + .stat(self.path.as_ref()) + .await + .map_err(error)? + .content_length(), + }; + Ok(size.saturating_sub(length)..size) + } + } + } +} +#[async_trait] +impl FileReader for Reader { + async fn read_range(&self, range: ReadRange) -> Result { + Ok(self + .inner + .read(self.range(range).await?) + .await + .map_err(error)? + .to_bytes()) + } + async fn read_ranges(&self, ranges: Vec>) -> Result> { + Ok(self + .inner + .fetch(ranges) + .await + .map_err(error)? + .into_iter() + .map(|b| b.to_bytes()) + .collect()) + } + fn stream( + self: Arc, + range: Option, + ) -> BoxStream<'static, Result> { + futures::stream::once(async move { + let stream = match range { + Some(range) => { + self.inner + .clone() + .into_bytes_stream(self.range(range).await?) + .await + } + None => self.inner.clone().into_bytes_stream(..).await, + } + .map_err(error)?; + Ok::<_, Error>(stream.map_err(Error::Io)) + }) + .try_flatten() + .boxed() + } +} + +fn directory_path(path: &Path) -> String { + if path.as_ref().is_empty() { + String::new() + } else { + format!("{path}/") + } +} diff --git a/datafusion/storage-opendal/tests/contracts.rs b/datafusion/storage-opendal/tests/contracts.rs new file mode 100644 index 0000000000000..238bed2a3ab4f --- /dev/null +++ b/datafusion/storage-opendal/tests/contracts.rs @@ -0,0 +1,26 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#[path = "../../storage/tests/common/backend.rs"] +mod backend; + +#[tokio::test] +async fn opendal_adapter_contract() { + let operator = opendal::Operator::new(opendal::services::Memory::default()).unwrap(); + backend::backend_contract(&datafusion_storage_opendal::OpendalStorage::new(operator)) + .await; +} diff --git a/datafusion/storage/Cargo.toml b/datafusion/storage/Cargo.toml new file mode 100644 index 0000000000000..10696dc06b240 --- /dev/null +++ b/datafusion/storage/Cargo.toml @@ -0,0 +1,47 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +[lints] +workspace = true + +[package] +name = "datafusion-storage" +description = "File access contracts and storage registration for DataFusion" +version.workspace = true +edition.workspace = true +license.workspace = true +rust-version.workspace = true +readme = "README.md" +authors.workspace = true +homepage.workspace = true +repository.workspace = true + +[dependencies] +async-trait.workspace = true +bytes.workspace = true +chrono.workspace = true +futures.workspace = true +parking_lot.workspace = true +percent-encoding.workspace = true +thiserror = "2" +tokio = { workspace = true, features = ["io-util", "sync"] } +tokio-util = { workspace = true, features = ["io", "rt"] } +url.workspace = true + +[dev-dependencies] +tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } diff --git a/datafusion/storage/README.md b/datafusion/storage/README.md new file mode 100644 index 0000000000000..4659ba07e2202 --- /dev/null +++ b/datafusion/storage/README.md @@ -0,0 +1,159 @@ + + +# DataFusion Storage + +The file stack uses two backend-independent traits: `Storage` for a namespace's +file operations and `FileReader` for one open file. Outputs implement Tokio's +existing `AsyncWrite` trait. SDK implementations live in +`datafusion-storage-object-store` and `datafusion-storage-opendal`. + +This PoC replaces the storage dependency boundary while preserving the operations +used by DataFusion's built-in formats. It explores +[issue #14854](https://github.com/apache/datafusion/issues/14854) together with the +backend opt-out in [PR #25144](https://github.com/apache/datafusion/pull/25144). + +## Register once + +Configure the SDK client, credentials, and middleware in the application, then +register it once: + +```rust,ignore +use datafusion::prelude::*; +use datafusion::storage::OpendalStorage; +use std::sync::Arc; +use url::Url; + +let ctx = SessionContext::new(); +ctx.register_storage( + &Url::parse("s3://warehouse")?, + Arc::new(OpendalStorage::new(operator)), +)?; +let df = ctx.read_parquet("s3://warehouse/events/", ParquetReadOptions::default()).await?; +df.filter(col("id").gt(lit(10)))?.show().await?; +``` + +ObjectStore uses the same registration method with +`Arc::new(ObjectStoreStorage::new(store))`. The adapter crates can be used directly; +DataFusion also re-exports them under its `object_store` and `opendal` features. + +The registration serves the ordinary CSV, JSON, Arrow, Avro, and Parquet APIs, +listing tables, SQL `COPY TO`, and file table `INSERT`. Output support remains +format-dependent. Sessions sharing a `RuntimeEnv` share registrations. Default +session assembly installs the local ObjectStore adapter when enabled, preserving +an existing application registration. A bare `RuntimeEnv` has an empty registry. + +## Interfaces and ownership + +```mermaid +flowchart TD + Context[SessionContext.register_storage] --> Registry[StorageRegistry] + Registry --> Binding[StorageBinding: namespace and registration identity] + Binding --> Storage[Storage trait] + Storage --> Metadata[stat / list / list_with_delimiter] + Storage --> Reader[open: FileReader] + Storage --> Output[writer: Tokio AsyncWrite] + Reader --> Reads[read_range / read_ranges / stream] + ObjectStore[ObjectStore adapter] -. implements .-> Storage + OpenDAL[OpenDAL adapter] -. implements .-> Storage + Formats[Built-in formats] --> Reader + Formats --> Output +``` + +`Storage` owns the client and its configuration. Its operations use paths relative +to that namespace. `open` creates an owned reader without requiring a preceding +HEAD or a known file size. `FileReader` keeps single-range and batched reads +separate, allowing adapters to retain native implementations of both. Suffix +reads serve Arrow IPC footers; full-file streams do not require metadata size. +Buffers retain their ownership independently of the reader. + +The registry normalizes scheme, host, and port. Paths are file locations within a +namespace, not registry mounts. Plans retain their immutable binding across +replacement or deregistration. Cache keys include the binding identity. Planning +and subsequent executions receive fresh `FileAccessContext` values; backends can +use these to associate I/O with queries. Dropping a scan cancels its context and +releases its pending futures, without promising that remote work is reversed. + +`FileInfo` holds path, size, modification time, ETag, and version independently of +SDK types. ETag and version remain observational metadata: the adapter does not +automatically enforce conditional reads or a snapshot. Format-specific decoding, +validation, pruning, and metadata caches remain in their existing format layers. +Custom `ParquetFileReaderFactory` implementations continue to take precedence. + +`StorageParquetTable::try_new(binding, files, schema)` accepts explicit `FileInfo` +values for manifest inputs. Supplying a schema skips inference. This path does +not call listing. A read-only backend implements `open`; unsupported operations +can use the trait's default errors. SDK adapters call their backend directly and +propagate errors instead of checking capabilities or selecting fallback providers. + +## Reads and writes + +ObjectStore single and batched ranges call `get_range` and `get_ranges` directly. +Full, bounded, and suffix streams use the corresponding `get_opts` requests. +Delimiter listing calls `list_with_delimiter` so listing tables can skip pruned +subdirectories. Existing SDK clients and wrappers remain in the request path. + +OpenDAL uses its native reader, batched fetch, and byte stream APIs. Its reader +accepts offset ranges, so the adapter resolves file size when a suffix read needs +it, using observed reader metadata when available. Other reads do not acquire a +size solely to satisfy the DataFusion interface. + +Writers return `Box` to existing encoders and +compression wrappers. `shutdown` completes the output through the backend. +`WriterOptions.buffer_size` is passed to ObjectStore's `BufWriter` capacity or +OpenDAL's writer chunk configuration. There is no additional universal output +buffer, conditional-create mode, output metadata lookup, or abort state machine. +Failure and drop retain the underlying writer's cleanup behavior. + +## Dependency boundary + +The storage contract has no SDK dependency. Execution, listing, and file-format +crates depend on this contract. Each adapter depends on the contract and its own +SDK; neither depends on the other adapter. Backend features only control adapter +re-exports and default assembly, not individual file operations. + +For example, `default-features = false, features = ["sql", "parquet", "avro", "opendal"]` +retains the selected file formats without a normal ObjectStore dependency. +SDK-based test fixtures can still introduce development dependencies. + +`register_storage` replaces `register_object_store`. There is one registry and +one set of ordinary format entry points. Public file metadata and paths no longer +use SDK types. Custom formats and reader factories receive `StorageBinding` and +DataFusion file types. + +## Limits and validation + +Avro and sequential Arrow IPC paths buffer full inputs, including local files. +This is the accepted consequence of removing the local SDK payload branch in +this PoC; other streaming formats keep incremental decoding. CPU `Bytes` do not +imply GPU memory or zero-copy guarantees. Cloud service behavior requires +service-specific validation. + +Adapter tests cover ordinary writes, single/batched/suffix reads, streams, +recursive and delimiter listing, and metadata preservation. Integration tests +cover the normal format APIs, COPY, INSERT, manifest inputs, binding replacement, +custom Parquet factories, and query context lifetimes. Dependency checks ensure +that the OpenDAL-only file stack does not pull in ObjectStore. + +```sh +cargo test -p datafusion-storage -p datafusion-storage-object-store -p datafusion-storage-opendal +cargo test -p datafusion --no-default-features --features sql,parquet,avro,opendal --test storage +cargo tree -p datafusion --no-default-features --features sql,parquet,avro,opendal --edges normal +cargo clippy --workspace --all-features --all-targets -- -D warnings +``` diff --git a/datafusion/storage/src/lib.rs b/datafusion/storage/src/lib.rs new file mode 100644 index 0000000000000..1687fa91ff036 --- /dev/null +++ b/datafusion/storage/src/lib.rs @@ -0,0 +1,177 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Backend-independent file operations and storage registration. + +pub mod path; +mod read; +mod registry; + +pub use read::{FileReader, ReadRange}; +pub use registry::{StorageBinding, StorageRegistry, StorageUrl}; + +use async_trait::async_trait; +use chrono::{DateTime, Utc}; +use futures::{StreamExt, stream::BoxStream}; +use path::Path; +use std::{fmt::Debug, sync::Arc}; +use tokio::io::AsyncWrite; +use tokio_util::sync::CancellationToken; + +/// An error independent of the selected storage SDK. +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("File not found: {0}")] + NotFound(String), + #[error("Storage operation not supported: {0}")] + NotSupported(String), + #[error("Invalid file access: {0}")] + InvalidInput(String), + #[error(transparent)] + Path(#[from] path::Error), + #[error(transparent)] + Url(#[from] url::ParseError), + #[error(transparent)] + Io(#[from] std::io::Error), + #[error("{backend}: {source}")] + Backend { + backend: &'static str, + #[source] + source: Box, + }, +} + +pub type Result = std::result::Result; + +/// File metadata relative to a storage namespace. +/// +/// ETags and versions are observations returned by the backend. Reading a file +/// does not automatically turn these fields into conditional requests. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct FileInfo { + pub location: Path, + pub size: u64, + pub last_modified: DateTime, + pub e_tag: Option, + pub version: Option, +} + +impl FileInfo { + /// Describe an explicitly supplied file without performing metadata I/O. + pub fn new(location: Path, size: u64) -> Self { + Self { + location, + size, + last_modified: DateTime::UNIX_EPOCH, + e_tag: None, + version: None, + } + } +} + +/// Access state scoped to one planning operation or query execution. +/// Backends may use it to associate their work with the calling query. +#[derive(Clone, Debug, Default)] +pub struct FileAccessContext { + pub query_id: Arc, + pub cancellation: CancellationToken, +} + +impl FileAccessContext { + pub fn new(query_id: impl Into>) -> Self { + Self { + query_id: query_id.into(), + cancellation: CancellationToken::new(), + } + } +} + +/// Immediate files and common directory prefixes returned by delimiter listing. +#[derive(Debug, Default)] +pub struct DirectoryListing { + pub files: Vec, + pub directories: Vec, +} + +/// Existing output buffering configuration, interpreted by the backend writer. +#[derive(Debug, Clone, Copy, Default)] +pub struct WriterOptions { + pub buffer_size: Option, +} + +/// Output consumed by the existing format encoders and compression wrappers. +/// `shutdown` completes the output using the backend's normal write semantics. +pub type FileOutput = Box; + +/// File operations within one storage namespace. +/// +/// Implementations own clients, configuration, and backend-specific behavior. +/// Unsupported operations fail at the backend; registration does not probe +/// capabilities or select a different backend. Default errors allow read-only +/// implementations to omit operations they do not provide. +#[async_trait] +pub trait Storage: Debug + Send + Sync { + /// Open a reusable reader without requiring a preceding metadata request. + async fn open( + &self, + path: &Path, + context: FileAccessContext, + ) -> Result>; + + async fn stat(&self, _path: &Path, _context: &FileAccessContext) -> Result { + Err(Error::NotSupported("stat".into())) + } + + /// Recursively enumerate files matching a path prefix. + fn list( + &self, + _prefix: &Path, + _context: FileAccessContext, + ) -> BoxStream<'_, Result> { + futures::stream::once(async { Err(Error::NotSupported("list".into())) }).boxed() + } + + /// List direct children, preserving backend directory-pruning support. + async fn list_with_delimiter( + &self, + _prefix: &Path, + _context: &FileAccessContext, + ) -> Result { + Err(Error::NotSupported("list_with_delimiter".into())) + } + + async fn writer( + &self, + _path: &Path, + _options: WriterOptions, + _context: FileAccessContext, + ) -> Result { + Err(Error::NotSupported("writer".into())) + } +} + +/// Collect a byte stream for formats that require a complete in-memory input. +pub async fn collect_bytes( + mut stream: BoxStream<'_, Result>, +) -> Result { + use futures::TryStreamExt; + let mut bytes = bytes::BytesMut::new(); + while let Some(chunk) = stream.try_next().await? { + bytes.extend_from_slice(&chunk); + } + Ok(bytes.freeze()) +} diff --git a/datafusion/storage/src/path/mod.rs b/datafusion/storage/src/path/mod.rs new file mode 100644 index 0000000000000..eef074e782474 --- /dev/null +++ b/datafusion/storage/src/path/mod.rs @@ -0,0 +1,815 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Backend-independent file paths. Adapted from Apache Arrow ObjectStore. + +use percent_encoding::percent_decode; +use std::fmt::Formatter; +#[cfg(not(target_arch = "wasm32"))] +use url::Url; + +/// The delimiter to separate object namespaces, creating a directory structure. +pub const DELIMITER: &str = "/"; + +/// The path delimiter as a single byte +pub const DELIMITER_BYTE: u8 = DELIMITER.as_bytes()[0]; + +/// The path delimiter as a single char +pub const DELIMITER_CHAR: char = DELIMITER_BYTE as char; + +mod parts; + +pub use parts::{InvalidPart, PathPart, PathParts}; + +/// Error returned by [`Path::parse`] +#[derive(Debug, thiserror::Error)] +#[non_exhaustive] +pub enum Error { + /// Error when there's an empty segment between two slashes `/` in the path + #[error("Path \"{}\" contained empty path segment", path)] + EmptySegment { + /// The source path + path: String, + }, + + /// Error when an invalid segment is encountered in the given path + #[error("Error parsing Path \"{}\": {}", path, source)] + BadSegment { + /// The source path + path: String, + /// The part containing the error + source: Box, + }, + + /// Error when path cannot be canonicalized + #[error("Failed to canonicalize path \"{}\": {}", path.display(), source)] + Canonicalize { + /// The source path + path: std::path::PathBuf, + /// The underlying error + source: std::io::Error, + }, + + /// Error when the path is not a valid URL + #[error("Unable to convert path \"{}\" to URL", path.display())] + InvalidPath { + /// The source path + path: std::path::PathBuf, + }, + + /// Error when a path contains non-unicode characters + #[error("Path \"{}\" contained non-unicode characters: {}", path, source)] + NonUnicode { + /// The source path + path: String, + /// The underlying `UTF8Error` + source: std::str::Utf8Error, + }, + + /// Error when the a path doesn't start with given prefix + #[error("Path {} does not start with prefix {}", path, prefix)] + PrefixMismatch { + /// The source path + path: String, + /// The mismatched prefix + prefix: String, + }, +} + +/// A parsed path representation that can be safely written to object storage +/// +/// A [`Path`] maintains the following invariants: +/// +/// * Paths are delimited by `/` +/// * Paths do not contain leading or trailing `/` +/// * Paths do not contain relative path segments, i.e. `.` or `..` +/// * Paths do not contain empty path segments +/// * Paths do not contain any ASCII control characters +/// +/// There are no enforced restrictions on path length, however, it should be noted that most +/// object stores do not permit paths longer than 1024 bytes, and many filesystems do not +/// support path segments longer than 255 bytes. +/// +/// # Encode +/// +/// In theory object stores support any UTF-8 character sequence, however, certain character +/// sequences cause compatibility problems with some applications and protocols. Additionally +/// some filesystems may impose character restrictions, see [`LocalFileSystem`]. As such the +/// naming guidelines for [S3], [GCS] and [Azure Blob Storage] all recommend sticking to a +/// limited character subset. +/// +/// [S3]: https://docs.aws.amazon.com/AmazonS3/latest/userguide/object-keys.html +/// [GCS]: https://cloud.google.com/storage/docs/naming-objects +/// [Azure Blob Storage]: https://docs.microsoft.com/en-us/rest/api/storageservices/Naming-and-Referencing-Containers--Blobs--and-Metadata#blob-names +/// +/// A string containing potentially problematic path segments can therefore be encoded to a [`Path`] +/// using [`Path::from`] or [`Path::from_iter`]. This will percent encode any problematic +/// segments according to [RFC 1738]. +/// +/// ``` +/// # use datafusion_storage::path::Path; +/// assert_eq!(Path::from("foo/bar").as_ref(), "foo/bar"); +/// assert_eq!(Path::from("foo//bar").as_ref(), "foo/bar"); +/// assert_eq!(Path::from("foo/../bar").as_ref(), "foo/%2E%2E/bar"); +/// assert_eq!(Path::from("/").as_ref(), ""); +/// assert_eq!(Path::from_iter(["foo", "foo/bar"]).as_ref(), "foo/foo%2Fbar"); +/// ``` +/// +/// Note: if provided with an already percent encoded string, this will encode it again +/// +/// ``` +/// # use datafusion_storage::path::Path; +/// assert_eq!(Path::from("foo/foo%2Fbar").as_ref(), "foo/foo%252Fbar"); +/// ``` +/// +/// # Parse +/// +/// Alternatively a [`Path`] can be parsed from an existing string, returning an +/// error if it is invalid. Unlike the encoding methods above, this will permit +/// arbitrary unicode, including percent encoded sequences. +/// +/// ``` +/// # use datafusion_storage::path::Path; +/// assert_eq!(Path::parse("/foo/foo%2Fbar").unwrap().as_ref(), "foo/foo%2Fbar"); +/// Path::parse("..").unwrap_err(); // Relative path segments are disallowed +/// Path::parse("/foo//").unwrap_err(); // Empty path segments are disallowed +/// Path::parse("\x00").unwrap_err(); // ASCII control characters are disallowed +/// ``` +/// +/// [RFC 1738]: https://www.ietf.org/rfc/rfc1738.txt +/// [`LocalFileSystem`]: https://docs.rs/object_store/latest/object_store/local/struct.LocalFileSystem.html +#[derive(Debug, Clone, Default, PartialEq, Eq, Hash, Ord, PartialOrd)] +pub struct Path { + /// The raw path with no leading or trailing delimiters + raw: String, +} + +impl Path { + /// Iterate over the validated path segments. + pub fn iter(&self) -> PathParts<'_> { + self.into_iter() + } + + /// An empty [`Path`] that points to the root of the store, equivalent to `Path::from("/")`. + /// + /// See also [`Path::is_root`]. + /// + /// # Example + /// + /// ``` + /// # use datafusion_storage::path::Path; + /// assert_eq!(Path::ROOT, Path::from("/")); + /// ``` + pub const ROOT: Self = Self { raw: String::new() }; + + /// Parse a string as a [`Path`], returning a [`Error`] if invalid, + /// as defined on the docstring for [`Path`] + /// + /// Note: this will strip any leading `/` or trailing `/` + pub fn parse(path: impl AsRef) -> Result { + let path = path.as_ref(); + + let stripped = path.strip_prefix(DELIMITER).unwrap_or(path); + if stripped.is_empty() { + return Ok(Default::default()); + } + + let stripped = stripped.strip_suffix(DELIMITER).unwrap_or(stripped); + + for segment in stripped.split(DELIMITER) { + if segment.is_empty() { + return Err(Error::EmptySegment { path: path.into() }); + } + + PathPart::parse(segment).map_err(|source| { + let path = path.into(); + Error::BadSegment { + source: Box::new(source), + path, + } + })?; + } + + Ok(Self { + raw: stripped.to_string(), + }) + } + + #[cfg(not(target_arch = "wasm32"))] + /// Convert a filesystem path to a [`Path`] relative to the filesystem root + /// + /// This will return an error if the path contains illegal character sequences + /// as defined on the docstring for [`Path`] or does not exist + /// + /// Note: this will canonicalize the provided path, resolving any symlinks + pub fn from_filesystem_path( + path: impl AsRef, + ) -> Result { + let absolute = std::fs::canonicalize(&path).map_err(|source| { + let path = path.as_ref().into(); + Error::Canonicalize { source, path } + })?; + + Self::from_absolute_path(absolute) + } + + #[cfg(not(target_arch = "wasm32"))] + /// Convert an absolute filesystem path to a [`Path`] relative to the filesystem root + /// + /// This will return an error if the path contains illegal character sequences, + /// as defined on the docstring for [`Path`], or `base` is not an absolute path + pub fn from_absolute_path(path: impl AsRef) -> Result { + Self::from_absolute_path_with_base(path, None) + } + + #[cfg(not(target_arch = "wasm32"))] + /// Convert a filesystem path to a [`Path`] relative to the provided base + /// + /// This will return an error if the path contains illegal character sequences, + /// as defined on the docstring for [`Path`], or `base` does not refer to a parent + /// path of `path`, or `base` is not an absolute path + pub(crate) fn from_absolute_path_with_base( + path: impl AsRef, + base: Option<&Url>, + ) -> Result { + let url = absolute_path_to_url(path)?; + let path = match base { + Some(prefix) => url.path().strip_prefix(prefix.path()).ok_or_else(|| { + Error::PrefixMismatch { + path: url.path().to_string(), + prefix: prefix.to_string(), + } + })?, + None => url.path(), + }; + + // Reverse any percent encoding performed by conversion to URL + Self::from_url_path(path) + } + + /// Parse a url encoded string as a [`Path`], returning a [`Error`] if invalid + /// + /// This will return an error if the path contains illegal character sequences + /// as defined on the docstring for [`Path`] + pub fn from_url_path(path: impl AsRef) -> Result { + let path = path.as_ref(); + let decoded = + percent_decode(path.as_bytes()) + .decode_utf8() + .map_err(|source| { + let path = path.into(); + Error::NonUnicode { source, path } + })?; + + Self::parse(decoded) + } + + /// Returns the number of [`PathPart`]s in this [`Path`] + /// + /// This is equivalent to calling `.parts().count()` manually. + /// + /// # Performance + /// + /// This operation is `O(n)`. + #[doc(alias = "len")] + pub fn parts_count(&self) -> usize { + self.raw.split_terminator(DELIMITER).count() + } + + /// True if this [`Path`] points to the root of the store, equivalent to `Path::from("/")`. + /// + /// See also [`Path::ROOT`]. + /// + /// # Example + /// + /// ``` + /// # use datafusion_storage::path::Path; + /// assert!(Path::from("/").is_root()); + /// assert!(Path::parse("").unwrap().is_root()); + /// ``` + pub fn is_root(&self) -> bool { + self.raw.is_empty() + } + + /// Returns the [`PathPart`]s of this [`Path`] + /// + /// Equivalent to calling `.into_iter()` on a `&Path`. + pub fn parts(&self) -> PathParts<'_> { + PathParts::new(&self.raw) + } + + /// Returns a copy of this [`Path`] with the last path segment removed + /// + /// Returns `None` if this path has zero segments. + pub fn parent(&self) -> Option { + if self.raw.is_empty() { + return None; + } + + let Some((prefix, _filename)) = self.raw.rsplit_once(DELIMITER) else { + return Some(Self::ROOT); + }; + + Some(Self { + raw: prefix.to_string(), + }) + } + + /// Returns the last path segment containing the filename stored in this [`Path`] + /// + /// Returns `None` only if this path is the root path. + pub fn filename(&self) -> Option<&str> { + match self.raw.is_empty() { + true => None, + false => self.raw.rsplit(DELIMITER).next(), + } + } + + /// Returns the extension of the file stored in this [`Path`], if any + pub fn extension(&self) -> Option<&str> { + self.filename() + .and_then(|f| f.rsplit_once('.')) + .and_then(|(_, extension)| { + if extension.is_empty() { + None + } else { + Some(extension) + } + }) + } + + /// Returns an iterator of the [`PathPart`] of this [`Path`] after `prefix` + /// + /// Returns `None` if the prefix does not match. + pub fn prefix_match( + &self, + prefix: &Self, + ) -> Option> + '_> { + let mut stripped = self.raw.strip_prefix(&prefix.raw)?; + if !stripped.is_empty() && !prefix.raw.is_empty() { + stripped = stripped.strip_prefix(DELIMITER)?; + } + Some(PathParts::new(stripped)) + } + + /// Returns true if this [`Path`] starts with `prefix` + pub fn prefix_matches(&self, prefix: &Self) -> bool { + self.prefix_match(prefix).is_some() + } + + /// Creates a new child of this [`Path`] + #[deprecated = "use .join() or .clone().join() instead"] + pub fn child<'a>(&self, child: impl Into>) -> Self { + self.clone().join(child) + } + + /// Appends a single path segment to this [`Path`] + pub fn join<'a>(self, child: impl Into>) -> Self { + let child_cow_str = child.into().raw; + + let raw = if self.raw.is_empty() { + child_cow_str.to_string() + } else { + use std::fmt::Write; + + let mut raw = self.raw; + write!(raw, "{DELIMITER}{child_cow_str}") + .expect("failed to append to string"); + raw + }; + + Self { raw } + } +} + +impl AsRef for Path { + fn as_ref(&self) -> &str { + &self.raw + } +} + +impl From<&str> for Path { + fn from(path: &str) -> Self { + Self::from_iter(path.split(DELIMITER)) + } +} + +impl From for Path { + fn from(path: String) -> Self { + Self::from_iter(path.split(DELIMITER)) + } +} + +impl From for String { + fn from(path: Path) -> Self { + path.raw + } +} + +impl std::fmt::Display for Path { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + self.raw.fmt(f) + } +} + +impl<'a, I> FromIterator for Path +where + I: Into>, +{ + fn from_iter>(iter: T) -> Self { + let mut this = Self::ROOT; + this.extend(iter); + this + } +} + +/// See also [`Path::parts`] +impl<'a> IntoIterator for &'a Path { + type Item = PathPart<'a>; + type IntoIter = PathParts<'a>; + + fn into_iter(self) -> Self::IntoIter { + PathParts::new(&self.raw) + } +} + +/// [`Path`] supports appending [`PathPart`]s of one `Path` to another `Path`. +/// +/// # Examples +/// +/// Suppose Alice is copying Bob's file to her own user directory. +/// We could choose the full path of the new file by taking the original +/// absolute path, making it relative to Bob's home +/// +/// ```rust +/// # use datafusion_storage::path::Path; +/// let alice_home = Path::from("Users/alice"); +/// let bob_home = Path::from("Users/bob"); +/// let bob_file = Path::from("Users/bob/documents/file.txt"); +/// +/// let mut alice_file = alice_home; +/// alice_file.extend(bob_file.prefix_match(&bob_home).unwrap()); +/// +/// assert_eq!(alice_file, Path::from("Users/alice/documents/file.txt")); +/// ``` +impl<'a, I: Into>> Extend for Path { + fn extend>(&mut self, iter: T) { + for s in iter { + let s = s.into(); + if !s.raw.is_empty() { + if !self.raw.is_empty() { + self.raw.push(DELIMITER_CHAR); + } + self.raw.push_str(&s.raw); + } + } + } +} + +#[cfg(not(target_arch = "wasm32"))] +/// Given an absolute filesystem path convert it to a URL representation without canonicalization +pub(crate) fn absolute_path_to_url( + path: impl AsRef, +) -> Result { + Url::from_file_path(&path).map_err(|_| Error::InvalidPath { + path: path.as_ref().into(), + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn delimiter_char_is_forward_slash() { + assert_eq!(DELIMITER_CHAR, '/'); + } + + #[test] + fn cloud_prefix_with_trailing_delimiter() { + // Use case: files exist in object storage named `foo/bar.json` and + // `foo_test.json`. A search for the prefix `foo/` should return + // `foo/bar.json` but not `foo_test.json'. + let prefix = Path::from_iter(["test"]); + assert_eq!(prefix.as_ref(), "test"); + } + + #[test] + fn push_encodes() { + let location = Path::from_iter(["foo/bar", "baz%2Ftest"]); + assert_eq!(location.as_ref(), "foo%2Fbar/baz%252Ftest"); + } + + #[test] + fn test_parse() { + assert_eq!(Path::parse("/").unwrap().as_ref(), ""); + assert_eq!(Path::parse("").unwrap().as_ref(), ""); + + let err = Path::parse("//").unwrap_err(); + assert!(matches!(err, Error::EmptySegment { .. })); + + assert_eq!(Path::parse("/foo/bar/").unwrap().as_ref(), "foo/bar"); + assert_eq!(Path::parse("foo/bar/").unwrap().as_ref(), "foo/bar"); + assert_eq!(Path::parse("foo/bar").unwrap().as_ref(), "foo/bar"); + + let err = Path::parse("foo///bar").unwrap_err(); + assert!(matches!(err, Error::EmptySegment { .. })); + } + + #[test] + fn convert_raw_before_partial_eq() { + // dir and file_name + let cloud = Path::from("test_dir/test_file.json"); + let built = Path::from_iter(["test_dir", "test_file.json"]); + + assert_eq!(built, cloud); + + // dir and file_name w/o dot + let cloud = Path::from("test_dir/test_file"); + let built = Path::from_iter(["test_dir", "test_file"]); + + assert_eq!(built, cloud); + + // dir, no file + let cloud = Path::from("test_dir/"); + let built = Path::from_iter(["test_dir"]); + assert_eq!(built, cloud); + + // file_name, no dir + let cloud = Path::from("test_file.json"); + let built = Path::from_iter(["test_file.json"]); + assert_eq!(built, cloud); + + // empty + let cloud = Path::from(""); + let built = Path::from_iter(["", ""]); + + assert_eq!(built, cloud); + } + + #[test] + fn parts_after_prefix_behavior() { + let existing_path = Path::from("apple/bear/cow/dog/egg.json"); + + // Prefix with one directory + let prefix = Path::from("apple"); + let expected_parts: Vec> = vec!["bear", "cow", "dog", "egg.json"] + .into_iter() + .map(Into::into) + .collect(); + let parts: Vec<_> = existing_path.prefix_match(&prefix).unwrap().collect(); + assert_eq!(parts, expected_parts); + + // Prefix with two directories + let prefix = Path::from("apple/bear"); + let expected_parts: Vec> = vec!["cow", "dog", "egg.json"] + .into_iter() + .map(Into::into) + .collect(); + let parts: Vec<_> = existing_path.prefix_match(&prefix).unwrap().collect(); + assert_eq!(parts, expected_parts); + + // Not a prefix + let prefix = Path::from("cow"); + assert!(existing_path.prefix_match(&prefix).is_none()); + + // Prefix with a partial directory + let prefix = Path::from("ap"); + assert!(existing_path.prefix_match(&prefix).is_none()); + + // Prefix matches but there aren't any parts after it + let existing = Path::from("apple/bear/cow/dog"); + + assert_eq!(existing.prefix_match(&existing).unwrap().count(), 0); + assert_eq!(Path::default().parts().count(), 0); + } + + #[test] + fn parts_count() { + assert_eq!(Path::ROOT.parts().count(), Path::ROOT.parts_count()); + + let path = path("foo/bar/baz"); + assert_eq!(path.parts_count(), 3); + assert_eq!(path.parts_count(), path.parts().count()); + } + + #[test] + fn prefix_matches_raw_content() { + assert_eq!(Path::ROOT.parent(), None, "empty path must have no prefix"); + + assert_eq!(path("foo").parent().unwrap(), Path::ROOT); + assert_eq!(path("foo/bar").parent().unwrap(), path("foo")); + assert_eq!(path("foo/bar/baz").parent().unwrap(), path("foo/bar")); + } + + #[test] + fn prefix_matches() { + let haystack = Path::from_iter(["foo/bar", "baz%2Ftest", "something"]); + // self starts with self + assert!( + haystack.prefix_matches(&haystack), + "{haystack:?} should have started with {haystack:?}" + ); + + // a longer prefix doesn't match + let needle = haystack.clone().join("longer now"); + assert!( + !haystack.prefix_matches(&needle), + "{haystack:?} shouldn't have started with {needle:?}" + ); + + // one dir prefix matches + let needle = Path::from_iter(["foo/bar"]); + assert!( + haystack.prefix_matches(&needle), + "{haystack:?} should have started with {needle:?}" + ); + + // two dir prefix matches + let needle = needle.join("baz%2Ftest"); + assert!( + haystack.prefix_matches(&needle), + "{haystack:?} should have started with {needle:?}" + ); + + // partial dir prefix doesn't match + let needle = Path::from_iter(["f"]); + assert!( + !haystack.prefix_matches(&needle), + "{haystack:?} should not have started with {needle:?}" + ); + + // one dir and one partial dir doesn't match + let needle = Path::from_iter(["foo/bar", "baz"]); + assert!( + !haystack.prefix_matches(&needle), + "{haystack:?} should not have started with {needle:?}" + ); + + // empty prefix matches + let needle = Path::from(""); + assert!( + haystack.prefix_matches(&needle), + "{haystack:?} should have started with {needle:?}" + ); + } + + #[test] + fn prefix_matches_with_file_name() { + let haystack = + Path::from_iter(["foo/bar", "baz%2Ftest", "something", "foo.segment"]); + + // All directories match and file name is a prefix + let needle = Path::from_iter(["foo/bar", "baz%2Ftest", "something", "foo"]); + + assert!( + !haystack.prefix_matches(&needle), + "{haystack:?} should not have started with {needle:?}" + ); + + // All directories match but file name is not a prefix + let needle = Path::from_iter(["foo/bar", "baz%2Ftest", "something", "e"]); + + assert!( + !haystack.prefix_matches(&needle), + "{haystack:?} should not have started with {needle:?}" + ); + + // Not all directories match; file name is a prefix of the next directory; this + // does not match + let needle = Path::from_iter(["foo/bar", "baz%2Ftest", "s"]); + + assert!( + !haystack.prefix_matches(&needle), + "{haystack:?} should not have started with {needle:?}" + ); + + // Not all directories match; file name is NOT a prefix of the next directory; + // no match + let needle = Path::from_iter(["foo/bar", "baz%2Ftest", "p"]); + + assert!( + !haystack.prefix_matches(&needle), + "{haystack:?} should not have started with {needle:?}" + ); + } + + #[test] + fn path_containing_spaces() { + let a = Path::from_iter(["foo bar", "baz"]); + let b = Path::from("foo bar/baz"); + let c = Path::parse("foo bar/baz").unwrap(); + + assert_eq!(a.raw, "foo bar/baz"); + assert_eq!(a.raw, b.raw); + assert_eq!(b.raw, c.raw); + } + + #[test] + fn from_url_path() { + let a = Path::from_url_path("foo%20bar").unwrap(); + let b = Path::from_url_path("foo/%2E%2E/bar").unwrap_err(); + let c = Path::from_url_path("foo%2F%252E%252E%2Fbar").unwrap(); + let d = Path::from_url_path("foo/%252E%252E/bar").unwrap(); + let e = Path::from_url_path("%48%45%4C%4C%4F").unwrap(); + let f = Path::from_url_path("foo/%FF/as").unwrap_err(); + + assert_eq!(a.raw, "foo bar"); + assert!(matches!(b, Error::BadSegment { .. })); + assert_eq!(c.raw, "foo/%2E%2E/bar"); + assert_eq!(d.raw, "foo/%2E%2E/bar"); + assert_eq!(e.raw, "HELLO"); + assert!(matches!(f, Error::NonUnicode { .. })); + } + + #[test] + fn filename_from_path() { + let a = Path::from("foo/bar"); + let b = Path::from("foo/bar.baz"); + let c = Path::from("foo.bar/baz"); + + assert_eq!(a.filename(), Some("bar")); + assert_eq!(b.filename(), Some("bar.baz")); + assert_eq!(c.filename(), Some("baz")); + } + + #[test] + fn file_extension() { + let a = Path::from("foo/bar"); + let b = Path::from("foo/bar.baz"); + let c = Path::from("foo.bar/baz"); + let d = Path::from("foo.bar/baz.qux"); + + assert_eq!(a.extension(), None); + assert_eq!(b.extension(), Some("baz")); + assert_eq!(c.extension(), None); + assert_eq!(d.extension(), Some("qux")); + } + + #[test] + fn root_is_root() { + assert!(Path::ROOT.is_root()); + assert!(Path::ROOT.parts().next().is_none()); + } + + /// Main test for `impl Extend for Path`, covers most cases. + #[test] + fn impl_extend() { + let mut p = Path::ROOT; + + p.extend(&Path::ROOT); + assert_eq!(p, Path::ROOT); + + p.extend(&path("foo")); + assert_eq!(p, path("foo")); + + p.extend(&path("bar/baz")); + assert_eq!(p, path("foo/bar/baz")); + + p.extend(&path("a/b/c")); + assert_eq!(p, path("foo/bar/baz/a/b/c")); + } + + /// Test for `impl Extend for Path`, specifically covers addition of a single segment. + #[test] + fn impl_extend_for_one_segment() { + let mut p = Path::ROOT; + + p.extend(&path("foo")); + assert_eq!(p, path("foo")); + + p.extend(&path("bar")); + assert_eq!(p, path("foo/bar")); + + p.extend(&path("baz")); + assert_eq!(p, path("foo/bar/baz")); + } + + #[test] + fn parent() { + assert_eq!(Path::ROOT.parent(), None); + assert_eq!(path("foo").parent(), Some(Path::ROOT)); + assert_eq!(path("foo/bar").parent(), Some(path("foo"))); + assert_eq!(path("foo/bar/baz").parent(), Some(path("foo/bar"))); + } + + /// Construct a [`Path`] from a raw `&str`, or panic trying. + #[track_caller] + fn path(raw: &str) -> Path { + Path::parse(raw).unwrap() + } +} diff --git a/datafusion/storage/src/path/parts.rs b/datafusion/storage/src/path/parts.rs new file mode 100644 index 0000000000000..05e43f6867970 --- /dev/null +++ b/datafusion/storage/src/path/parts.rs @@ -0,0 +1,208 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use percent_encoding::{AsciiSet, CONTROLS, percent_encode}; +use std::{ + borrow::Cow, + iter::{self, FusedIterator}, + str::SplitTerminator, +}; + +use crate::path::DELIMITER_BYTE; + +/// Error returned by [`PathPart::parse`] +#[derive(Debug, thiserror::Error)] +#[error( + "Encountered illegal character sequence \"{}\" whilst parsing path segment \"{}\"", + illegal, + segment +)] +pub struct InvalidPart { + segment: String, + illegal: String, +} + +/// The PathPart type exists to validate the directory/file names that form part +/// of a path. +/// +/// A [`PathPart`] is guaranteed to: +/// +/// * Contain no ASCII control characters or `/` +/// * Not be a relative path segment, i.e. `.` or `..` +#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Debug, Default, Hash)] +pub struct PathPart<'a> { + pub(super) raw: Cow<'a, str>, +} + +impl<'a> PathPart<'a> { + /// Parse the provided path segment as a [`PathPart`] returning an error if invalid + pub fn parse(segment: &'a str) -> Result { + if segment == "." || segment == ".." { + return Err(InvalidPart { + segment: segment.to_string(), + illegal: segment.to_string(), + }); + } + + for c in segment.chars() { + if c.is_ascii_control() || c == '/' { + return Err(InvalidPart { + segment: segment.to_string(), + // This is correct as only single byte characters up to this point + illegal: c.to_string(), + }); + } + } + + Ok(Self { + raw: segment.into(), + }) + } +} + +/// Characters we want to encode. +const INVALID: &AsciiSet = &CONTROLS + // The delimiter we are reserving for internal hierarchy + .add(DELIMITER_BYTE) + // Characters AWS recommends avoiding for object keys + // https://docs.aws.amazon.com/AmazonS3/latest/dev/UsingMetadata.html + .add(b'\\') + .add(b'{') + .add(b'^') + .add(b'}') + .add(b'%') + .add(b'`') + .add(b']') + .add(b'"') // " <-- my editor is confused about double quotes within single quotes + .add(b'>') + .add(b'[') + .add(b'~') + .add(b'<') + .add(b'#') + .add(b'|') + // Characters Google Cloud Storage recommends avoiding for object names + // https://cloud.google.com/storage/docs/naming-objects + .add(b'\r') + .add(b'\n') + .add(b'*') + .add(b'?'); + +impl<'a> From<&'a [u8]> for PathPart<'a> { + fn from(v: &'a [u8]) -> Self { + let inner = match v { + // We don't want to encode `.` generally, but we do want to disallow parts of paths + // to be equal to `.` or `..` to prevent file system traversal shenanigans. + b"." => "%2E".into(), + b".." => "%2E%2E".into(), + other => percent_encode(other, INVALID).into(), + }; + Self { raw: inner } + } +} + +impl<'a> From<&'a str> for PathPart<'a> { + fn from(v: &'a str) -> Self { + Self::from(v.as_bytes()) + } +} + +impl From for PathPart<'static> { + fn from(s: String) -> Self { + Self { + raw: Cow::Owned(PathPart::from(s.as_str()).raw.into_owned()), + } + } +} + +impl AsRef for PathPart<'_> { + fn as_ref(&self) -> &str { + self.raw.as_ref() + } +} + +/// See [`Path::parts`](super::Path::parts) +#[derive(Debug, Clone)] +pub struct PathParts<'a>(iter::Map, fn(&str) -> PathPart<'_>>); + +impl<'a> PathParts<'a> { + /// Create an iterator over the parts of the provided raw [`Path`](super::Path). + pub(super) fn new(raw: &'a str) -> Self { + Self( + raw.split_terminator(super::DELIMITER_CHAR) + .map(|s| PathPart { raw: s.into() }), + ) + } +} + +impl<'a> Iterator for PathParts<'a> { + type Item = PathPart<'a>; + + fn next(&mut self) -> Option { + self.0.next() + } +} + +impl FusedIterator for PathParts<'_> {} + +impl DoubleEndedIterator for PathParts<'_> { + fn next_back(&mut self) -> Option { + self.0.next_back() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn path_part_delimiter_gets_encoded() { + let part: PathPart<'_> = "foo/bar".into(); + assert_eq!(part.raw, "foo%2Fbar"); + } + + #[test] + fn path_part_given_already_encoded_string() { + let part: PathPart<'_> = "foo%2Fbar".into(); + assert_eq!(part.raw, "foo%252Fbar"); + } + + #[test] + fn path_part_cant_be_one_dot() { + let part: PathPart<'_> = ".".into(); + assert_eq!(part.raw, "%2E"); + } + + #[test] + fn path_part_cant_be_two_dots() { + let part: PathPart<'_> = "..".into(); + assert_eq!(part.raw, "%2E%2E"); + } + + #[test] + fn path_part_parse() { + PathPart::parse("foo").unwrap(); + PathPart::parse("foo/bar").unwrap_err(); + + // Test percent-encoded path + PathPart::parse("foo%2Fbar").unwrap(); + PathPart::parse("L%3ABC.parquet").unwrap(); + + // Test path containing bad escape sequence + PathPart::parse("%Z").unwrap(); + PathPart::parse("%%").unwrap(); + } +} diff --git a/datafusion/storage/src/read.rs b/datafusion/storage/src/read.rs new file mode 100644 index 0000000000000..01b04ff0ffe98 --- /dev/null +++ b/datafusion/storage/src/read.rs @@ -0,0 +1,53 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::Result; +use async_trait::async_trait; +use bytes::Bytes; +use futures::stream::BoxStream; +use std::{fmt::Debug, ops::Range, sync::Arc}; + +/// A bounded range or a suffix read, as used by existing file formats. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ReadRange { + Bounded(Range), + Suffix(u64), +} + +impl From> for ReadRange { + fn from(range: Range) -> Self { + Self::Bounded(range) + } +} + +/// A reusable reader for one file. +/// +/// Single and batched reads remain separate so adapters can preserve native +/// implementations of both. Batch results correspond to input ranges in order. +/// Returned buffers own their bytes independently of this reader. Range and EOF +/// errors follow the backend; this interface does not enforce file revisions or +/// validate reads against previously observed metadata. +#[async_trait] +pub trait FileReader: Debug + Send + Sync { + async fn read_range(&self, range: ReadRange) -> Result; + async fn read_ranges(&self, ranges: Vec>) -> Result>; + /// Stream a range, or the full file when no range is supplied. + fn stream( + self: Arc, + range: Option, + ) -> BoxStream<'static, Result>; +} diff --git a/datafusion/storage/src/registry.rs b/datafusion/storage/src/registry.rs new file mode 100644 index 0000000000000..5d663e687f237 --- /dev/null +++ b/datafusion/storage/src/registry.rs @@ -0,0 +1,171 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::{Error, Result, Storage}; +use parking_lot::RwLock; +use std::{ + collections::HashMap, + fmt::{Display, Formatter}, + sync::Arc, +}; +use url::{Position, Url}; + +/// A storage namespace, normalized independently of the chosen backend. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct StorageUrl(Url); +impl StorageUrl { + pub fn parse(value: impl AsRef) -> Result { + Self::new(&Url::parse(value.as_ref())?) + } + pub fn new(url: &Url) -> Result { + let key = format!( + "{}://{}", + url.scheme(), + &url[Position::BeforeHost..Position::AfterPort] + ); + let mut url = Url::parse(&key)?; + url.set_path("/"); + Ok(Self(url)) + } + pub fn local_filesystem() -> Self { + Self::parse("file://").unwrap() + } + pub fn as_str(&self) -> &str { + self.0.as_str() + } +} +impl AsRef for StorageUrl { + fn as_ref(&self) -> &Url { + &self.0 + } +} +impl Display for StorageUrl { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + self.0.fmt(f) + } +} + +/// A single authoritative map of storage registrations. +#[derive(Debug, Default)] +pub struct StorageRegistry { + storages: RwLock>>, +} +impl StorageRegistry { + pub fn register( + &self, + url: &Url, + storage: Arc, + ) -> Result>> { + let key = StorageUrl::new(url)?; + let storage = Arc::new(StorageBinding::new(key.clone(), storage)); + Ok(self.storages.write().insert(key, storage)) + } + /// Install a default without replacing an application registration. + pub fn register_default(&self, url: &Url, storage: Arc) -> Result<()> { + let key = StorageUrl::new(url)?; + self.storages + .write() + .entry(key.clone()) + .or_insert_with(|| Arc::new(StorageBinding::new(key, storage))); + Ok(()) + } + pub fn get(&self, url: &Url) -> Result> { + let key = StorageUrl::new(url)?; + self.storages + .read() + .get(&key) + .cloned() + .ok_or_else(|| Error::NotFound(format!("storage registration {key}"))) + } + pub fn deregister(&self, url: &Url) -> Result> { + let key = StorageUrl::new(url)?; + self.storages + .write() + .remove(&key) + .ok_or_else(|| Error::NotFound(format!("storage registration {key}"))) + } +} + +impl AsRef for StorageUrl { + fn as_ref(&self) -> &str { + self.as_str() + } +} + +/// An immutable registration: a namespace and its backend. +/// Plans retain this binding across replacement or deregistration. +#[derive(Debug)] +pub struct StorageBinding { + id: u64, + url: StorageUrl, + storage: Arc, +} +impl StorageBinding { + pub fn new(url: StorageUrl, storage: Arc) -> Self { + static NEXT_ID: std::sync::atomic::AtomicU64 = + std::sync::atomic::AtomicU64::new(1); + let id = NEXT_ID.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + assert_ne!(id, 0, "storage binding identity exhausted"); + Self { id, url, storage } + } + /// Identity of this registration, independent of its replaceable URL. + pub fn id(&self) -> u64 { + self.id + } + pub fn url(&self) -> &StorageUrl { + &self.url + } + pub fn storage(&self) -> &Arc { + &self.storage + } + pub async fn open( + &self, + file: &crate::FileInfo, + context: crate::FileAccessContext, + ) -> Result> { + self.storage.open(&file.location, context).await + } + pub async fn stat( + &self, + path: &crate::path::Path, + context: &crate::FileAccessContext, + ) -> Result { + self.storage.stat(path, context).await + } + pub fn list( + &self, + path: &crate::path::Path, + context: crate::FileAccessContext, + ) -> futures::stream::BoxStream<'_, Result> { + self.storage.list(path, context) + } + pub async fn list_with_delimiter( + &self, + path: &crate::path::Path, + context: &crate::FileAccessContext, + ) -> Result { + self.storage.list_with_delimiter(path, context).await + } + pub async fn writer( + &self, + path: &crate::path::Path, + options: crate::WriterOptions, + context: crate::FileAccessContext, + ) -> Result { + self.storage.writer(path, options, context).await + } +} diff --git a/datafusion/storage/tests/common/backend.rs b/datafusion/storage/tests/common/backend.rs new file mode 100644 index 0000000000000..b1986eccc8aa1 --- /dev/null +++ b/datafusion/storage/tests/common/backend.rs @@ -0,0 +1,97 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use bytes::Bytes; +use datafusion_storage::{path::Path, *}; +use futures::TryStreamExt; +use tokio::io::AsyncWriteExt; + +pub async fn backend_contract(storage: &dyn Storage) { + let context = FileAccessContext::new("backend-contract"); + let path = Path::from("data.parquet"); + let mut writer = storage + .writer( + &path, + WriterOptions { + buffer_size: Some(8), + }, + context.clone(), + ) + .await + .unwrap(); + writer.write_all(b"01234").await.unwrap(); + writer.flush().await.unwrap(); + writer.write_all(b"56789").await.unwrap(); + writer.shutdown().await.unwrap(); + assert_eq!(storage.stat(&path, &context).await.unwrap().size, 10); + let reader = storage.open(&path, context.clone()).await.unwrap(); + assert_eq!(reader.read_range((2..5).into()).await.unwrap(), b"234"[..]); + assert_eq!( + reader.read_range(ReadRange::Suffix(3)).await.unwrap(), + b"789"[..] + ); + assert_eq!( + reader + .read_ranges(vec![7..10, 1..5, 1..5, 4..8]) + .await + .unwrap(), + vec![ + Bytes::from_static(b"789"), + Bytes::from_static(b"1234"), + Bytes::from_static(b"1234"), + Bytes::from_static(b"4567"), + ] + ); + assert_eq!( + collect_bytes(reader.clone().stream(Some((2..8).into()))) + .await + .unwrap(), + b"234567"[..] + ); + assert_eq!( + collect_bytes(reader.stream(None)).await.unwrap(), + b"0123456789"[..] + ); + for name in [ + "events/part.parquet", + "events/day/part.parquet", + "events-other/part.parquet", + ] { + let mut writer = storage + .writer(&Path::from(name), WriterOptions::default(), context.clone()) + .await + .unwrap(); + writer.write_all(b"data").await.unwrap(); + writer.shutdown().await.unwrap(); + } + let files = storage + .list(&Path::from("events"), context.clone()) + .try_collect::>() + .await + .unwrap(); + assert_eq!(files.len(), 2); + let children = storage + .list_with_delimiter(&Path::from("events"), &context) + .await + .unwrap(); + assert_eq!(children.files.len(), 1); + assert_eq!( + children.files[0].location, + Path::from("events/part.parquet") + ); + assert_eq!(children.directories, vec![Path::from("events/day")]); +} diff --git a/datafusion/storage/tests/contracts.rs b/datafusion/storage/tests/contracts.rs new file mode 100644 index 0000000000000..9b5cd962adb28 --- /dev/null +++ b/datafusion/storage/tests/contracts.rs @@ -0,0 +1,134 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use async_trait::async_trait; +use bytes::Bytes; +use datafusion_storage::{path::Path, *}; +use futures::{StreamExt, stream::BoxStream}; +use std::{ops::Range, sync::Arc}; + +#[derive(Debug)] +struct Memory(Bytes); +#[async_trait] +impl Storage for Memory { + async fn open(&self, _: &Path, _: FileAccessContext) -> Result> { + Ok(Arc::new(Self(self.0.clone()))) + } +} +#[async_trait] +impl FileReader for Memory { + async fn read_range(&self, range: ReadRange) -> Result { + let range = match range { + ReadRange::Bounded(range) => range, + ReadRange::Suffix(size) => { + (self.0.len() as u64).saturating_sub(size)..self.0.len() as u64 + } + }; + Ok(self.0.slice(range.start as usize..range.end as usize)) + } + async fn read_ranges(&self, ranges: Vec>) -> Result> { + Ok(ranges + .into_iter() + .map(|r| self.0.slice(r.start as usize..r.end as usize)) + .collect()) + } + fn stream( + self: Arc, + range: Option, + ) -> BoxStream<'static, Result> { + futures::stream::once(async move { + match range { + Some(range) => self.read_range(range).await, + None => Ok(self.0.clone()), + } + }) + .boxed() + } +} + +#[tokio::test] +async fn replacement_and_deregistration_preserve_bound_readers() { + let registry = StorageRegistry::default(); + let url = url::Url::parse("memory://bucket/prefix").unwrap(); + registry + .register(&url, Arc::new(Memory(Bytes::from_static(b"old")))) + .unwrap(); + let old = registry.get(&url).unwrap(); + let replaced = registry + .register(&url, Arc::new(Memory(Bytes::from_static(b"new")))) + .unwrap() + .unwrap(); + assert!(Arc::ptr_eq(&old, &replaced)); + let current = registry.get(&url).unwrap(); + assert_ne!(old.id(), current.id()); + registry.deregister(&url).unwrap(); + assert!(registry.get(&url).is_err()); + for (binding, expected) in [(old, b"old"), (current, b"new")] { + let reader = binding + .storage() + .open(&Path::from("file"), FileAccessContext::default()) + .await + .unwrap(); + let bytes = reader.read_range((0..3).into()).await.unwrap(); + drop(reader); + drop(binding); + assert_eq!(bytes, expected[..]); + } +} + +#[tokio::test] +async fn read_only_storage_does_not_require_metadata_or_output_operations() { + let storage = Memory(Bytes::from_static(b"data")); + let path = Path::from("file"); + let context = FileAccessContext::default(); + let reader = storage.open(&path, context.clone()).await.unwrap(); + assert_eq!( + collect_bytes(reader.stream(None)).await.unwrap(), + b"data"[..] + ); + assert!(matches!( + storage.stat(&path, &context).await, + Err(Error::NotSupported(_)) + )); + assert!(matches!( + storage + .writer(&path, WriterOptions::default(), context) + .await, + Err(Error::NotSupported(_)) + )); +} + +#[test] +fn urls_are_namespaces_not_path_mounts() { + assert_eq!(StorageUrl::local_filesystem().as_str(), "file:///"); + assert_eq!( + StorageUrl::parse("s3://username:password@host:123/foo?bar=baz") + .unwrap() + .as_str(), + "s3://host:123/", + ); + assert!(StorageUrl::parse("s3://bucket:invalid").is_err()); + + assert_eq!( + StorageUrl::parse("s3://bucket/a").unwrap(), + StorageUrl::parse("s3://bucket/b").unwrap() + ); + assert_ne!( + StorageUrl::parse("https://example.com:8443/a").unwrap(), + StorageUrl::parse("https://example.com/a").unwrap() + ); +} diff --git a/datafusion/substrait/Cargo.toml b/datafusion/substrait/Cargo.toml index a0f203cec8db6..1efa595265519 100644 --- a/datafusion/substrait/Cargo.toml +++ b/datafusion/substrait/Cargo.toml @@ -40,15 +40,16 @@ chrono = { workspace = true } datafusion = { workspace = true, features = ["sql"] } half = { workspace = true } itertools = { workspace = true } -object_store = { workspace = true } # We need to match the version in substrait, so we don't use the workspace version here pbjson-types = { version = "0.8.0" } prost = { workspace = true } substrait = { version = "0.63.0", features = ["serde"] } url = { workspace = true } -tokio = { workspace = true, features = ["fs"] } +tokio = { workspace = true, features = ["fs", "io-util"] } [dev-dependencies] +datafusion-storage = { workspace = true } +object_store = { workspace = true } datafusion = { workspace = true, features = ["nested_expressions", "unicode_expressions"] } datafusion-functions-aggregate = { workspace = true } serde_json = { workspace = true } diff --git a/datafusion/substrait/src/physical_plan/consumer.rs b/datafusion/substrait/src/physical_plan/consumer.rs index 3c0df01b4ad38..24f2082ab5e8f 100644 --- a/datafusion/substrait/src/physical_plan/consumer.rs +++ b/datafusion/substrait/src/physical_plan/consumer.rs @@ -21,13 +21,13 @@ use std::sync::Arc; use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::common::{not_impl_err, substrait_err}; use datafusion::datasource::listing::PartitionedFile; -use datafusion::datasource::object_store::ObjectStoreUrl; use datafusion::datasource::physical_plan::{ FileGroup, FileScanConfigBuilder, ParquetSource, }; use datafusion::error::{DataFusionError, Result}; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::SessionContext; +use datafusion::storage::StorageUrl; use crate::variation_const::{ DEFAULT_CONTAINER_TYPE_VARIATION_REF, LARGE_CONTAINER_TYPE_VARIATION_REF, @@ -36,7 +36,7 @@ use crate::variation_const::{ use async_recursion::async_recursion; use chrono::DateTime; use datafusion::datasource::memory::DataSourceExec; -use object_store::ObjectMeta; +use datafusion::storage::FileInfo; use substrait::proto::Type; use substrait::proto::read_rel::local_files::file_or_files::PathType; use substrait::proto::r#type::{Kind, Nullability}; @@ -82,7 +82,7 @@ pub async fn from_substrait_rel( let schema = Arc::new(Schema::new(fields)); let source = Arc::new(ParquetSource::new(Arc::clone(&schema))); base_config_builder = FileScanConfigBuilder::new( - ObjectStoreUrl::local_filesystem(), + StorageUrl::local_filesystem(), source, ); } @@ -107,7 +107,7 @@ pub async fn from_substrait_rel( )) }?; - // TODO substrait plans do not have `last_modified` or `size` but `ObjectMeta` + // TODO substrait plans do not have `last_modified` or `size` but `FileInfo` // requires them both - perhaps we can change the object-store crate // to make these optional? We cannot guarantee that we have access to the // files to get this information, depending on how this library is being @@ -119,14 +119,13 @@ pub async fn from_substrait_rel( .unwrap(); let size = 0; - let partitioned_file = - PartitionedFile::new_from_meta(ObjectMeta { - last_modified: last_modified.into(), - location: path.into(), - size, - e_tag: None, - version: None, - }); + let partitioned_file = PartitionedFile::new_from_meta(FileInfo { + last_modified: last_modified.into(), + location: path.into(), + size, + e_tag: None, + version: None, + }); let part_index = file.partition_index as usize; while part_index >= file_groups.len() { diff --git a/datafusion/substrait/tests/cases/roundtrip_physical_plan.rs b/datafusion/substrait/tests/cases/roundtrip_physical_plan.rs index 9773cf4aba10f..c28ec033c5a8f 100644 --- a/datafusion/substrait/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/substrait/tests/cases/roundtrip_physical_plan.rs @@ -21,13 +21,13 @@ use std::sync::Arc; use datafusion::arrow::datatypes::Schema; use datafusion::dataframe::DataFrame; use datafusion::datasource::listing::PartitionedFile; -use datafusion::datasource::object_store::ObjectStoreUrl; use datafusion::datasource::physical_plan::{ FileGroup, FileScanConfigBuilder, ParquetSource, }; use datafusion::error::Result; use datafusion::physical_plan::{ExecutionPlan, displayable}; use datafusion::prelude::{ParquetReadOptions, SessionContext}; +use datafusion::storage::StorageUrl; use datafusion_substrait::physical_plan::{consumer, producer}; use datafusion::datasource::memory::DataSourceExec; @@ -38,19 +38,18 @@ async fn parquet_exec() -> Result<()> { let schema = Arc::new(Schema::empty()); let source = Arc::new(ParquetSource::new(schema.clone())); - let scan_config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source) - .with_file_groups(vec![ - FileGroup::new(vec![PartitionedFile::new( - "file://foo/part-0.parquet".to_string(), - 123, - )]), - FileGroup::new(vec![PartitionedFile::new( - "file://foo/part-1.parquet".to_string(), - 123, - )]), - ]) - .build(); + let scan_config = FileScanConfigBuilder::new(StorageUrl::local_filesystem(), source) + .with_file_groups(vec![ + FileGroup::new(vec![PartitionedFile::new( + "file://foo/part-0.parquet".to_string(), + 123, + )]), + FileGroup::new(vec![PartitionedFile::new( + "file://foo/part-1.parquet".to_string(), + 123, + )]), + ]) + .build(); let parquet_exec: Arc = DataSourceExec::from_data_source(scan_config); diff --git a/datafusion/wasmtest/Cargo.toml b/datafusion/wasmtest/Cargo.toml index d6cea2e68d384..2e8ef58bbc1db 100644 --- a/datafusion/wasmtest/Cargo.toml +++ b/datafusion/wasmtest/Cargo.toml @@ -60,6 +60,7 @@ getrandom = { version = "0.3", features = ["wasm_js"] } wasm-bindgen = "0.2.99" [dev-dependencies] +datafusion-storage-object-store = { workspace = true } bytes = { workspace = true } futures = { workspace = true } object_store = { workspace = true } diff --git a/datafusion/wasmtest/src/lib.rs b/datafusion/wasmtest/src/lib.rs index d8da5d4b4f323..980c28c2b0b81 100644 --- a/datafusion/wasmtest/src/lib.rs +++ b/datafusion/wasmtest/src/lib.rs @@ -269,7 +269,14 @@ mod test { store.put(&path, buffer.into()).await.unwrap(); let url = Url::parse("memory://").unwrap(); - session_ctx.register_object_store(&url, Arc::new(store)); + session_ctx + .register_storage( + &url, + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + Arc::new(store), + )), + ) + .unwrap(); session_ctx .register_parquet("a", "memory:///a.parquet", Default::default()) .await @@ -312,7 +319,13 @@ mod test { let url = Url::parse("memory://").unwrap(); let ctx = SessionContext::new(); - ctx.register_object_store(&url, Arc::new(store)); + ctx.register_storage( + &url, + Arc::new(datafusion_storage_object_store::ObjectStoreStorage::new( + Arc::new(store), + )), + ) + .unwrap(); let csv_options = CsvReadOptions::new() .has_header(true) diff --git a/docs/source/user-guide/crate-configuration.md b/docs/source/user-guide/crate-configuration.md index 09c65107e58c8..32456d6cc0339 100644 --- a/docs/source/user-guide/crate-configuration.md +++ b/docs/source/user-guide/crate-configuration.md @@ -52,6 +52,52 @@ datafusion = { git = "https://github.com/apache/datafusion", branch = "main", de More on [Cargo dependencies](https://doc.rust-lang.org/cargo/reference/specifying-dependencies.html#specifying-dependencies) +## Building without object_store + +The default-enabled `object_store` feature selects the ObjectStore adapter and +local file access installed during session construction. File formats, listing +tables, SQL `COPY TO`, and `INSERT` depend on backend-independent storage traits. +To use OpenDAL with the same APIs: + +```toml +datafusion = { version = "55.0.0", default-features = false, features = ["sql", "parquet", "avro", "opendal"] } +``` + +Register `OpendalStorage::new(operator)` once with +`SessionContext::register_storage`, then use the ordinary `read_parquet`, +`read_csv`, `read_json`, `read_arrow`, and `read_avro` APIs. Applications configure +OpenDAL services and layers through their own dependency. Custom `FileAccess` +implementations need neither SDK; manifest-driven input can omit discovery. +The `parquet_encryption` feature does not require ObjectStore. + +SQL planning, custom table providers, in-memory tables, execution, memory +management, and local disk spilling also work without either storage adapter. +Without a backend feature, register an implementation explicitly for file access. + +Cargo unifies features across dependencies: any dependency that enables +`object_store` brings the adapter back into the graph. Check normal dependencies +with `cargo tree --edges normal`; a lockfile entry alone does not mean the +optional dependency is enabled. SDK-based test fixtures may add development +dependencies independently of the production graph. + +### Migrating existing builds + +Replace `register_object_store(url, store)` with +`register_storage(url, ObjectStoreStorage::new(store))`. The wrapper is re-exported +under `datafusion::storage` with the `object_store` feature. There is one storage +registry; runtime and session initialization do not maintain parallel registrations. + +Shared file types now belong to `datafusion-storage`: `PartitionedFile` contains +`FileInfo`, file and cache paths use `datafusion_storage::path::Path`, and scan +configuration uses `StorageUrl`. SDK adapters provide explicit conversions for +ObjectStore metadata and paths. Custom cache and encryption factory +implementations must update their signatures to the independent types. + +The Storage registry is shared by sessions using the same runtime. It is +independent of the legacy ObjectStore registry; registrations are not mirrored. +See the [storage prototype scope](https://github.com/apache/datafusion/blob/main/datafusion/storage/README.md) +for supported discovery behavior and remaining migration work. + ## Optimizing Builds Here are several suggestions to get the Rust compiler to produce faster code when diff --git a/test-utils/Cargo.toml b/test-utils/Cargo.toml index 525c892a406f7..433f75f55a31d 100644 --- a/test-utils/Cargo.toml +++ b/test-utils/Cargo.toml @@ -33,5 +33,8 @@ workspace = true arrow = { workspace = true } chrono-tz = { version = "0.10.4", default-features = false } datafusion-common = { workspace = true, default-features = true } +datafusion-storage = { workspace = true } +datafusion-storage-object-store = { workspace = true } env_logger = { workspace = true } +object_store = { workspace = true } rand = { workspace = true } diff --git a/test-utils/src/lib.rs b/test-utils/src/lib.rs index 55717c717c4af..11f89870b0325 100644 --- a/test-utils/src/lib.rs +++ b/test-utils/src/lib.rs @@ -144,3 +144,5 @@ fn primary_key(schema: &Schema, column_names: &[&str]) -> Constraint { Constraint::PrimaryKey(indices) } + +pub mod storage; diff --git a/test-utils/src/storage.rs b/test-utils/src/storage.rs new file mode 100644 index 0000000000000..5e11925b72701 --- /dev/null +++ b/test-utils/src/storage.rs @@ -0,0 +1,36 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Storage fixtures shared by format and execution tests. +use datafusion_storage::{StorageBinding, StorageUrl}; +use datafusion_storage_object_store::ObjectStoreStorage; +use std::sync::Arc; + +/// Bind a fixture backend to the local namespace used by format tests. +pub fn object_store(store: Arc) -> Arc { + Arc::new(StorageBinding::new( + StorageUrl::local_filesystem(), + Arc::new(ObjectStoreStorage::new(store)), + )) +} + +pub fn local() -> Arc { + object_store(Arc::new(object_store::local::LocalFileSystem::new())) +} + +/// Convert SDK fixture metadata at the adapter boundary. +pub use datafusion_storage_object_store::file_info;