diff --git a/Cargo.lock b/Cargo.lock index ac4934d7..e8d2e62c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -365,6 +365,28 @@ dependencies = [ "zstd-safe", ] +[[package]] +name = "async-generic" +version = "1.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddf3728566eefa873833159754f5732fb0951d3649e6e5b891cc70d56dd41673" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "async-lock" +version = "3.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5fd03604047cee9b6ce9de9f70c6cd540a0520c813cbd49bae61f33ab80ed1dc" +dependencies = [ + "event-listener", + "event-listener-strategy", + "pin-project-lite", +] + [[package]] name = "async-stream" version = "0.3.6" @@ -428,6 +450,17 @@ version = "1.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" +[[package]] +name = "auto_impl" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ffdcb70bdbc4d478427380519163274ac86e52916e10f0a8889adf0f96d3fee7" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "autocfg" version = "1.5.0" @@ -592,6 +625,28 @@ dependencies = [ "zstd", ] +[[package]] +name = "beacon-arrow-zarr" +version = "0.1.0" +dependencies = [ + "arrow", + "crossbeam", + "futures", + "hifitime", + "indexmap 2.10.0", + "nd-arrow-array", + "object_store 0.12.3", + "parking_lot", + "pin-project", + "regex", + "serde_json", + "tokio", + "tracing", + "zarrs", + "zarrs_object_store", + "zarrs_storage", +] + [[package]] name = "beacon-common" version = "1.3.0" @@ -660,6 +715,7 @@ version = "1.3.0" dependencies = [ "arrow", "async-trait", + "beacon-arrow-zarr", "beacon-config", "beacon-formats", "chrono", @@ -693,12 +749,16 @@ dependencies = [ "async_zip", "beacon-arrow-netcdf", "beacon-arrow-odv", + "beacon-arrow-zarr", "beacon-common", + "chrono", "datafusion", "futures", "geoarrow", "geoarrow-array", "geoparquet", + "glob", + "nd-arrow-array", "object_store 0.12.3", "parking_lot", "parquet", @@ -709,6 +769,10 @@ dependencies = [ "tokio-util", "tracing", "typetag", + "url", + "zarrs", + "zarrs_object_store", + "zarrs_storage", ] [[package]] @@ -725,7 +789,7 @@ dependencies = [ "gsw", "inventory", "lazy_static", - "lru", + "lru 0.14.0", "ordered-float 5.0.0", "serde", "wkt 0.12.0", @@ -829,6 +893,19 @@ dependencies = [ "generic-array", ] +[[package]] +name = "blosc-src" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb68d27ab5ceb94ae9cd343f6fbc7bb84543496d547ed7c0db6718175fd41cb6" +dependencies = [ + "cc", + "libz-sys", + "lz4-sys", + "snappy_src", + "zstd-sys", +] + [[package]] name = "brotli" version = "8.0.2" @@ -861,6 +938,20 @@ name = "bytemuck" version = "1.24.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1fbdf580320f38b612e485521afda1ee26d10cc9884efaaa750d383e13e3c5f4" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] [[package]] name = "byteorder" @@ -976,6 +1067,15 @@ dependencies = [ "unicode-width", ] +[[package]] +name = "concurrent-queue" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ca0197aee26d1ae37445ee532fefce43251d24cc7c166799f4d46817f1d3973" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "const-random" version = "0.1.18" @@ -1051,6 +1151,15 @@ version = "2.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "19d374276b40fb8bbdee95aef7c7fa6b5316ec764510eb64b8dd0e2ed0d7e7f5" +[[package]] +name = "crc32c" +version = "0.6.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a47af21622d091a8f0fb295b88bc886ac74efcc613efc19f5d0b21de5c89e47" +dependencies = [ + "rustc_version", +] + [[package]] name = "crc32fast" version = "1.5.0" @@ -1060,6 +1169,28 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "crossbeam" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1137cd7e7fc0fb5d3c5a8678be38ec56e819125d8d7907411fe24ccb943faca8" +dependencies = [ + "crossbeam-channel", + "crossbeam-deque", + "crossbeam-epoch", + "crossbeam-queue", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-channel" +version = "0.5.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "82b8f8f868b36967f9606790d1903570de9ceaf870a7bf9fbbd3016d636a2cb2" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "crossbeam-deque" version = "0.8.6" @@ -1079,6 +1210,15 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "crossbeam-queue" +version = "0.3.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f58bbc28f91df819d0aa2a2c00cd19754769c2fad90579b3592b1c9ba7a3115" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "crossbeam-utils" version = "0.8.21" @@ -1832,6 +1972,27 @@ dependencies = [ "syn", ] +[[package]] +name = "derive_more" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "093242cf7570c207c83073cf82f79706fe7b8317e98620a47d5be7c3d8497678" +dependencies = [ + "derive_more-impl", +] + +[[package]] +name = "derive_more-impl" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bda628edc44c4bb645fbe0f758797143e4e07926f7ebf4e9bdfbd3d2ce621df3" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "unicode-xid", +] + [[package]] name = "digest" version = "0.10.7" @@ -1936,6 +2097,27 @@ dependencies = [ "windows-sys 0.60.2", ] +[[package]] +name = "event-listener" +version = "5.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e13b66accf52311f30a0db42147dadea9850cb48cd070028831ae5f5d4b856ab" +dependencies = [ + "concurrent-queue", + "parking", + "pin-project-lite", +] + +[[package]] +name = "event-listener-strategy" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8be9f3dfaaffdae2972880079a491a1a8bb7cbed0b8dd7a347f668b4150a3b93" +dependencies = [ + "event-listener", + "pin-project-lite", +] + [[package]] name = "fastrand" version = "2.3.0" @@ -2339,6 +2521,7 @@ version = "2.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "459196ed295495a68f7d7fe1d84f6c4b7ff0e21fe3017b2f283c6fac3ad803c9" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "num-traits", @@ -2380,6 +2563,12 @@ dependencies = [ "foldhash", ] +[[package]] +name = "hashbrown" +version = "0.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5419bdc4f6a9207fbeba6d11b604d481addf78ecd10c11ad51e76c2f6482748d" + [[package]] name = "hdf5-metno-sys" version = "0.9.1" @@ -3032,6 +3221,15 @@ dependencies = [ "vcpkg", ] +[[package]] +name = "link-cplusplus" +version = "1.0.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f78c730aaa7d0b9336a299029ea49f9ee53b0ed06e9202e8cb7db9bae7b8c82" +dependencies = [ + "cc", +] + [[package]] name = "linux-raw-sys" version = "0.9.4" @@ -3070,12 +3268,31 @@ dependencies = [ "hashbrown 0.15.5", ] +[[package]] +name = "lru" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bfe949189f46fabb938b3a9a0be30fdd93fd8a09260da863399a8cf3db756ec8" +dependencies = [ + "hashbrown 0.15.5", +] + [[package]] name = "lru-slab" version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" +[[package]] +name = "lz4-sys" +version = "1.11.1+lz4-1.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6bd8c0d6c6ed0cd30b3652886bb8711dc4bb01d637a68105a3d5158039b418e6" +dependencies = [ + "cc", + "libc", +] + [[package]] name = "lz4_flex" version = "0.11.5" @@ -3192,6 +3409,49 @@ dependencies = [ "windows-sys 0.59.0", ] +[[package]] +name = "moka" +version = "0.12.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8261cd88c312e0004c1d51baad2980c66528dfdb2bee62003e643a4d8f86b077" +dependencies = [ + "async-lock", + "crossbeam-channel", + "crossbeam-epoch", + "crossbeam-utils", + "equivalent", + "event-listener", + "futures-util", + "parking_lot", + "portable-atomic", + "rustc_version", + "smallvec", + "tagptr", + "uuid", +] + +[[package]] +name = "monostate" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3341a273f6c9d5bef1908f17b7267bbab0e95c9bf69a0d4dcf8e9e1b2c76ef67" +dependencies = [ + "monostate-impl", + "serde", + "serde_core", +] + +[[package]] +name = "monostate-impl" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4db6d5580af57bf992f59068d4ea26fd518574ff48d7639b255a36f9de6e7e9" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "nd-arrow-array" version = "1.2.2" @@ -3299,6 +3559,7 @@ version = "0.4.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" dependencies = [ + "bytemuck", "num-traits", ] @@ -3502,6 +3763,16 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b15813163c1d831bf4a13c3610c05c0d03b39feb07f7e09fa234dac9b15aaf39" +[[package]] +name = "page_size" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30d5b2194ed13191c1999ae0704b7839fb18384fa22e49b57eeaa97d79ce40da" +dependencies = [ + "libc", + "winapi", +] + [[package]] name = "parking" version = "2.2.1" @@ -3780,6 +4051,18 @@ dependencies = [ "serde", ] +[[package]] +name = "quick_cache" +version = "0.6.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba15f5bccfb18c666351668b97bbff66da5093f96757ca15299e4e594fe1316e" +dependencies = [ + "ahash", + "equivalent", + "hashbrown 0.16.0", + "parking_lot", +] + [[package]] name = "quinn" version = "0.11.8" @@ -3905,6 +4188,15 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "rayon_iter_concurrent_limit" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d09ee01023de07fa073ce14c37cbe0a9e099c6b0b60a29cf4af6d04d9553fed7" +dependencies = [ + "rayon", +] + [[package]] name = "recursive" version = "0.1.1" @@ -4248,18 +4540,28 @@ checksum = "1bc711410fbe7399f390ca1c3b60ad0f53f80e95c5eb935e52268a0e2cd49acc" [[package]] name = "serde" -version = "1.0.200" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.228" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ddc6f9cc94d67c0e21aaf7eda3a010fd3af78ebf6e096aa6e2e13c79749cce4f" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" dependencies = [ "serde_derive", ] [[package]] name = "serde_derive" -version = "1.0.200" +version = "1.0.228" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "856f046b9400cee3c8c94ed572ecdb752444c24528c035cd35882aad6f492bcb" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" dependencies = [ "proc-macro2", "quote", @@ -4272,6 +4574,7 @@ version = "1.0.120" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4e0d21c9a8cae1235ad58a00c11cb40d4b1e5c784f1ef2c537876ed6ffd8b7c5" dependencies = [ + "indexmap 2.10.0", "itoa", "ryu", "serde", @@ -4287,6 +4590,17 @@ dependencies = [ "serde", ] +[[package]] +name = "serde_repr" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "175ee3e80ae9982737ca543e96133087cbd9a485eecc3bc4de9c1a37b47ea59c" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "serde_urlencoded" version = "0.7.1" @@ -4433,6 +4747,16 @@ version = "1.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1b6b67fb9a61334225b5b790716f609cd58395f895b3fe8b328786812a40bc3b" +[[package]] +name = "snappy_src" +version = "0.2.5+snappy.1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4e1432067a55bcfb1fd522d2aca6537a4fcea32bba87ea86921226d14f9bad53" +dependencies = [ + "cc", + "link-cplusplus", +] + [[package]] name = "socket2" version = "0.5.10" @@ -4608,6 +4932,12 @@ dependencies = [ "windows", ] +[[package]] +name = "tagptr" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b2093cf4c8eb1e67749a6762251bc9cd836b6fc171623bd0a9d324d37af2417" + [[package]] name = "tar" version = "0.4.44" @@ -5030,6 +5360,18 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4a1a07cc7db3810833284e8d372ccdc6da29741639ecc70c9ec107df0fa6154c" +[[package]] +name = "unicode-xid" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" + +[[package]] +name = "unsafe_cell_slice" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6659959f702dcdaad77bd6e42a9409a32ceccc06943ec93c8a4306be00eb6cf1" + [[package]] name = "untrusted" version = "0.9.0" @@ -5768,6 +6110,163 @@ dependencies = [ "synstructure", ] +[[package]] +name = "zarrs" +version = "0.22.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad12c7c2b91d2f6871f21efc28fd5a302809d7246974fdd99ef55bd3b16b78a0" +dependencies = [ + "async-generic", + "async-lock", + "async-trait", + "blosc-src", + "bytemuck", + "bytes", + "crc32c", + "derive_more", + "flate2", + "futures", + "getrandom 0.3.3", + "half", + "inventory", + "itertools 0.14.0", + "itoa", + "lru 0.16.1", + "moka", + "ndarray", + "num", + "num-complex", + "quick_cache", + "rayon", + "rayon_iter_concurrent_limit", + "serde", + "serde_json", + "thiserror 2.0.15", + "thread_local", + "unsafe_cell_slice", + "uuid", + "zarrs_data_type", + "zarrs_filesystem", + "zarrs_metadata", + "zarrs_metadata_ext", + "zarrs_plugin", + "zarrs_registry", + "zarrs_storage", + "zstd", +] + +[[package]] +name = "zarrs_data_type" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e68a3b9e663cf4933afcd90f460ee72986fdf6c2b4d43d0441ad049b802342" +dependencies = [ + "derive_more", + "half", + "inventory", + "num", + "thiserror 2.0.15", + "zarrs_metadata", + "zarrs_plugin", +] + +[[package]] +name = "zarrs_filesystem" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e135c32621a3a5796d917768d5c7aa7f58be9480ae00778956b82ec6409150b" +dependencies = [ + "bytes", + "derive_more", + "itertools 0.14.0", + "libc", + "page_size", + "parking_lot", + "pathdiff", + "thiserror 2.0.15", + "walkdir", + "zarrs_storage", +] + +[[package]] +name = "zarrs_metadata" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "708b938e5af9e6564d7135fb2d1e05c0deff3d7124694ff6822aa01614a6c991" +dependencies = [ + "derive_more", + "half", + "monostate", + "serde", + "serde_json", + "thiserror 2.0.15", +] + +[[package]] +name = "zarrs_metadata_ext" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d4fb56ca32761b64c4b2a3db1097fbd29adfb321a129279b1db99be0a61d361a" +dependencies = [ + "derive_more", + "half", + "monostate", + "num", + "serde", + "serde_json", + "serde_repr", + "thiserror 2.0.15", + "zarrs_metadata", + "zarrs_registry", +] + +[[package]] +name = "zarrs_object_store" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a5d9d0d3db426dd50dfb0b5d7cc1660bda368caeb5cd8645c60c46bc4f261a19" +dependencies = [ + "async-trait", + "futures", + "object_store 0.12.3", + "zarrs_storage", +] + +[[package]] +name = "zarrs_plugin" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3c9e0514d4c50f44d11285d5df70e4e586486a39826579c9d87ddc3f3dac561" +dependencies = [ + "thiserror 2.0.15", +] + +[[package]] +name = "zarrs_registry" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebe4e55522eeb87eefab89017bef78cb823f86861fd8a3cc12e9f6538c348d57" +dependencies = [ + "regex", +] + +[[package]] +name = "zarrs_storage" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bc1037a8fa8c44ccb8f5c6c85753a63ddf296fb43280f28150f0f29fda8d301" +dependencies = [ + "async-trait", + "auto_impl", + "bytes", + "derive_more", + "futures", + "itertools 0.14.0", + "parking_lot", + "thiserror 2.0.15", + "unsafe_cell_slice", +] + [[package]] name = "zerocopy" version = "0.8.26" diff --git a/Cargo.toml b/Cargo.toml index 2e056246..5888abb0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,5 +1,5 @@ [workspace] -members = ["beacon-api", "beacon-arrow-netcdf", "beacon-arrow-odv", "beacon-common", "beacon-config", "beacon-core", "beacon-functions", "beacon-query", "beacon-planner", "beacon-data-lake", "beacon-formats"] +members = ["beacon-api", "beacon-arrow-netcdf", "beacon-arrow-odv", "beacon-common", "beacon-config", "beacon-core", "beacon-functions", "beacon-query", "beacon-planner", "beacon-data-lake", "beacon-formats", "beacon-arrow-zarr"] [workspace.dependencies] tokio = { version = "1.47.1", features = ["full"] } @@ -16,19 +16,19 @@ utoipa-axum = "0.2.0" utoipa-scalar = { version = "0.3.0", features = ["axum"] } utoipa-swagger-ui = { version = "9.0.0", features = ["axum"] } -serde = { version = "=1.0.200", features = ["rc", "derive"] } +serde = { version = "^1.0.2", features = ["rc", "derive"] } serde_json = "=1.0.120" anyhow = "1.0.95" thiserror = "2.0.12" tracing = "0.1.41" tracing-subscriber = { version = "0.3.19", features = ["env-filter"] } - glob = "0.3.2" tempfile = "3.15.0" typetag = "0.2.19" indexmap = { version = "2.7.1", features = ["serde"]} chrono = { version = "0.4.41", features = ["serde"] } +crossbeam = "0.8.4" datafusion = "49.0.0" object_store = { version = "0.12.3", features = ["aws"] } @@ -39,4 +39,3 @@ parquet = { version = "^55.2.0", features = ["async"] } geoarrow = { version = "=0.4.0" } geoarrow-array = "=0.4.0" geoparquet = "0.4.0" -pprof = { version = "0.15", features = ["flamegraph"] } diff --git a/beacon-api/Cargo.toml b/beacon-api/Cargo.toml index fb3aa485..ae914052 100644 --- a/beacon-api/Cargo.toml +++ b/beacon-api/Cargo.toml @@ -5,6 +5,7 @@ edition = "2021" [target.'cfg(not(windows))'.dependencies] tikv-jemallocator = "0.6.0" +pprof = { version = "0.15", features = ["flamegraph"] } [dependencies] axum = { version = "0.8.1", features = ["tracing"] } @@ -27,7 +28,6 @@ arrow-schema = { workspace = true} uuid = { version = "1.16.0" } tracing = { workspace = true} tracing-subscriber = { workspace = true} -pprof = { workspace = true } # Local dependencies beacon-config = { path = "../beacon-config" } diff --git a/beacon-arrow-netcdf/.DS_Store b/beacon-arrow-netcdf/.DS_Store new file mode 100644 index 00000000..8c9b71d3 Binary files /dev/null and b/beacon-arrow-netcdf/.DS_Store differ diff --git a/beacon-arrow-zarr/Cargo.toml b/beacon-arrow-zarr/Cargo.toml new file mode 100644 index 00000000..3e14c58c --- /dev/null +++ b/beacon-arrow-zarr/Cargo.toml @@ -0,0 +1,22 @@ +[package] +name = "beacon-arrow-zarr" +version = "0.1.0" +edition = "2024" + +[dependencies] +zarrs = {"version" = "0.22.2", features = ["async"]} +zarrs_object_store = "0.5.0" +zarrs_storage = "0.4.0" +object_store = { workspace = true } +tokio = { workspace = true } +arrow = { workspace = true } +nd-arrow-array = { git = "https://github.com/maris-development/nd-arrow-array.git", branch = "main", version = "1.2.2" } +serde_json = { workspace = true } +indexmap = { workspace = true } +tracing = { workspace = true } +hifitime = "4.0.2" +regex = "1.11.1" +futures = { workspace = true } +pin-project = "1.1.10" +parking_lot = { workspace = true } +crossbeam = { workspace = true } \ No newline at end of file diff --git a/beacon-arrow-zarr/src/array_slice_pushdown.rs b/beacon-arrow-zarr/src/array_slice_pushdown.rs new file mode 100644 index 00000000..48c92476 --- /dev/null +++ b/beacon-arrow-zarr/src/array_slice_pushdown.rs @@ -0,0 +1,60 @@ +#[derive(Clone, Debug)] +pub struct ArraySlicePushDown { + pub dimension: String, + pub start: Option, + pub end: Option, +} + +impl ArraySlicePushDown { + pub fn new(dimension: String, start: Option, end: Option) -> Self { + Self { + dimension, + start, + end, + } + } + + pub fn dimension(&self) -> &str { + &self.dimension + } + + pub fn start(&self) -> Option { + self.start + } + + pub fn end(&self) -> Option { + self.end + } + + pub fn overlapping_range( + &self, + array_range: std::ops::Range, + ) -> Option> { + match (self.start, self.end) { + (Some(push_start), Some(push_end)) => { + let push_end = push_end + 1; // Make end inclusive + let start = array_range.start.max(push_start); + let end = array_range.end.min(push_end); + if start < end { Some(start..end) } else { None } + } + (Some(push_start), None) => { + let start = array_range.start.max(push_start); + if start < array_range.end { + Some(start..array_range.end) + } else { + None + } + } + (None, Some(push_end)) => { + let push_end = push_end + 1; // Make end inclusive + let end = array_range.end.min(push_end); + if array_range.start < end { + Some(array_range.start..end) + } else { + None + } + } + (None, None) => Some(array_range), + } + } +} diff --git a/beacon-arrow-zarr/src/attributes.rs b/beacon-arrow-zarr/src/attributes.rs new file mode 100644 index 00000000..67db2aa1 --- /dev/null +++ b/beacon-arrow-zarr/src/attributes.rs @@ -0,0 +1,78 @@ +use std::sync::Arc; + +use nd_arrow_array::NdArrowArray; + +pub enum AttributeValue { + String(String), + Float64(f64), + Bool(bool), +} + +impl AttributeValue { + pub fn from_json_value(value: &serde_json::Value) -> Option { + match value { + serde_json::Value::String(s) => Some(AttributeValue::String(s.clone())), + serde_json::Value::Number(n) => n.as_f64().map(AttributeValue::Float64), + serde_json::Value::Bool(b) => Some(AttributeValue::Bool(*b)), + _ => None, + } + } + + pub fn as_str(&self) -> Option<&str> { + match self { + AttributeValue::String(s) => Some(s.as_str()), + _ => None, + } + } + + pub fn as_f64(&self) -> Option { + match self { + AttributeValue::Float64(f) => Some(*f), + _ => None, + } + } + + pub fn as_bool(&self) -> Option { + match self { + AttributeValue::Bool(b) => Some(*b), + _ => None, + } + } + + pub fn arrow_data_type(&self) -> arrow::datatypes::DataType { + match self { + AttributeValue::String(_) => arrow::datatypes::DataType::Utf8, + AttributeValue::Float64(_) => arrow::datatypes::DataType::Float64, + AttributeValue::Bool(_) => arrow::datatypes::DataType::Boolean, + } + } + + pub fn as_nd_arrow_array(&self) -> NdArrowArray { + match self { + AttributeValue::String(s) => { + let array = arrow::array::StringArray::from(vec![s.as_str()]); + NdArrowArray::new( + Arc::new(array) as arrow::array::ArrayRef, + nd_arrow_array::dimensions::Dimensions::Scalar, + ) + .unwrap() + } + AttributeValue::Float64(f) => { + let array = arrow::array::Float64Array::from(vec![*f]); + NdArrowArray::new( + Arc::new(array) as arrow::array::ArrayRef, + nd_arrow_array::dimensions::Dimensions::Scalar, + ) + .unwrap() + } + AttributeValue::Bool(b) => { + let array = arrow::array::BooleanArray::from(vec![*b]); + NdArrowArray::new( + Arc::new(array) as arrow::array::ArrayRef, + nd_arrow_array::dimensions::Dimensions::Scalar, + ) + .unwrap() + } + } + } +} diff --git a/beacon-arrow-zarr/src/data_types.rs b/beacon-arrow-zarr/src/data_types.rs new file mode 100644 index 00000000..ee5911b4 --- /dev/null +++ b/beacon-arrow-zarr/src/data_types.rs @@ -0,0 +1,20 @@ +pub fn try_zarrs_dtype_to_arrow( + data_type: zarrs::array::DataType, +) -> Result { + match data_type { + zarrs::array::DataType::Bool => Ok(arrow::datatypes::DataType::Boolean), + zarrs::array::DataType::Int8 => Ok(arrow::datatypes::DataType::Int8), + zarrs::array::DataType::Int16 => Ok(arrow::datatypes::DataType::Int16), + zarrs::array::DataType::Int32 => Ok(arrow::datatypes::DataType::Int32), + zarrs::array::DataType::Int64 => Ok(arrow::datatypes::DataType::Int64), + zarrs::array::DataType::UInt8 => Ok(arrow::datatypes::DataType::UInt8), + zarrs::array::DataType::UInt16 => Ok(arrow::datatypes::DataType::UInt16), + zarrs::array::DataType::UInt32 => Ok(arrow::datatypes::DataType::UInt32), + zarrs::array::DataType::UInt64 => Ok(arrow::datatypes::DataType::UInt64), + zarrs::array::DataType::Float32 => Ok(arrow::datatypes::DataType::Float32), + zarrs::array::DataType::Float64 => Ok(arrow::datatypes::DataType::Float64), + zarrs::array::DataType::String => Ok(arrow::datatypes::DataType::Utf8), + zarrs::array::DataType::Bytes => Ok(arrow::datatypes::DataType::Binary), + _ => Err(format!("Unsupported Zarrs data type: {:?}", data_type)), + } +} diff --git a/beacon-arrow-zarr/src/decoder.rs b/beacon-arrow-zarr/src/decoder.rs new file mode 100644 index 00000000..b870a836 --- /dev/null +++ b/beacon-arrow-zarr/src/decoder.rs @@ -0,0 +1,589 @@ +use std::{str::FromStr, sync::Arc}; + +use arrow::{ + array::{ + ArrayRef, AsArray, BinaryArray, BooleanArray, PrimitiveArray, StringArray, + TimestampMillisecondArray, + }, + datatypes::{ + Float32Type, Float64Type, Int8Type, Int16Type, Int32Type, Int64Type, SchemaRef, + TimestampMillisecondType, UInt8Type, UInt16Type, UInt32Type, UInt64Type, + }, +}; +use hifitime::Epoch; +use indexmap::IndexMap; +use nd_arrow_array::{NdArrowArray, batch::NdRecordBatch}; +use regex::Regex; + +use crate::attributes::AttributeValue; + +pub trait Decoder: Send + Sync { + fn create( + group_reader: &crate::reader::AsyncArrowZarrGroupReader, + input_schema: SchemaRef, + ) -> Result + where + Self: Sized; + fn decoded_schema(&self) -> SchemaRef; + fn decoding_array_pipeline( + &self, + array_name: &str, + ) -> Option>>; + fn decode_array(&self, array_name: &str, array: NdArrowArray) -> Result { + if let Some(pipeline) = self.decoding_array_pipeline(array_name) { + pipeline.iter().fold(Ok(array), |maybe_array, step| { + maybe_array.and_then(|array| step.decode(array)) + }) + } else { + Ok(array) + } + } +} + +pub trait DecodingPipelineStep: Send + Sync { + fn output_data_type( + &self, + input_data_type: arrow::datatypes::DataType, + ) -> arrow::datatypes::DataType { + input_data_type + } + fn decode(&self, array: NdArrowArray) -> Result; +} + +pub struct FillValueDecoder { + // Array decoders to apply per array name + array_fill_decoders: std::collections::HashMap>, + decoded_schema: SchemaRef, +} + +impl Decoder for FillValueDecoder { + fn create( + group_reader: &crate::reader::AsyncArrowZarrGroupReader, + input_schema: SchemaRef, + ) -> Result + where + Self: Sized, + { + let mut array_fill_decoders = std::collections::HashMap::new(); + for (array_name, array_reader) in group_reader.arrays() { + // Add a FillValueDecodingStep if the array has a fill value + let fill_value_attr = array_reader.attributes().get("_FillValue"); + if let Some(arrow_fill_value) = ArrowFillValue::from_zarrs_fill_value( + array_reader.data_type(), + fill_value_attr + .and_then(AttributeValue::from_json_value) + .as_ref(), + ) { + let fill_value_decoder = FillValueDecodingStep { + fill_value: arrow_fill_value, + }; + array_fill_decoders.insert( + array_name.to_string(), + Arc::new(fill_value_decoder) as Arc, + ); + } + } + Ok(Self { + array_fill_decoders, + decoded_schema: input_schema, + }) + } + + fn decoded_schema(&self) -> SchemaRef { + self.decoded_schema.clone() + } + + fn decoding_array_pipeline( + &self, + array_name: &str, + ) -> Option>> { + self.array_fill_decoders + .get(array_name) + .map(|decoder| vec![decoder.clone()]) + } +} + +pub struct FillValueDecodingStep { + fill_value: ArrowFillValue, +} + +pub enum ArrowFillValue { + Bool(bool), + Int8(i8), + Int16(i16), + Int32(i32), + Int64(i64), + UInt8(u8), + UInt16(u16), + UInt32(u32), + UInt64(u64), + Float32(f32), + Float64(f64), + String(String), + Bytes(Vec), +} + +impl ArrowFillValue { + pub fn from_zarrs_fill_value( + array_data_type: &zarrs::array::DataType, + fill_value_attr: Option<&AttributeValue>, + ) -> Option { + match array_data_type { + zarrs::array::DataType::Bool => fill_value_attr + .and_then(|attr| attr.as_bool()) + .map(ArrowFillValue::Bool), + zarrs::array::DataType::Int8 => fill_value_attr + .and_then(|attr| attr.as_f64().and_then(|v| i8::try_from(v as i64).ok())) + .map(ArrowFillValue::Int8), + zarrs::array::DataType::Int16 => fill_value_attr + .and_then(|attr| attr.as_f64().and_then(|v| i16::try_from(v as i64).ok())) + .map(ArrowFillValue::Int16), + zarrs::array::DataType::Int32 => fill_value_attr + .and_then(|attr| attr.as_f64().and_then(|v| i32::try_from(v as i64).ok())) + .map(ArrowFillValue::Int32), + zarrs::array::DataType::Int64 => fill_value_attr + .and_then(|attr| attr.as_f64().map(|v| v as i64).or(None)) + .map(ArrowFillValue::Int64), + zarrs::array::DataType::UInt8 => fill_value_attr + .and_then(|attr| attr.as_f64().and_then(|v| u8::try_from(v as u64).ok())) + .map(ArrowFillValue::UInt8), + zarrs::array::DataType::UInt16 => fill_value_attr + .and_then(|attr| attr.as_f64().and_then(|v| u16::try_from(v as u64).ok())) + .map(ArrowFillValue::UInt16), + zarrs::array::DataType::UInt32 => fill_value_attr + .and_then(|attr| attr.as_f64().and_then(|v| u32::try_from(v as u64).ok())) + .map(ArrowFillValue::UInt32), + + zarrs::array::DataType::UInt64 => fill_value_attr + .and_then(|attr| attr.as_f64().map(|v| v as u64).or(None)) + .map(ArrowFillValue::UInt64), + + zarrs::array::DataType::Float32 => fill_value_attr + .and_then(|attr| attr.as_f64().map(|v| v as f32).or(None)) + .map(ArrowFillValue::Float32), + + zarrs::array::DataType::Float64 => fill_value_attr + .and_then(|attr| attr.as_f64()) + .map(ArrowFillValue::Float64), + + zarrs::array::DataType::String => { + // Interpret fill_bytes as UTF-8 string + fill_value_attr + .and_then(|attr| attr.as_str().map(|s| s.to_string())) + .map(ArrowFillValue::String) + } + + _ => None, + } + } +} + +impl DecodingPipelineStep for FillValueDecodingStep { + fn output_data_type( + &self, + input_data_type: arrow::datatypes::DataType, + ) -> arrow::datatypes::DataType { + input_data_type + } + + fn decode(&self, array: NdArrowArray) -> Result { + let arrow_array = array.as_arrow_array(); + let nullable_array: ArrayRef = match (arrow_array.data_type(), &self.fill_value) { + (arrow::datatypes::DataType::Null, _) => return Ok(array.clone()), + (arrow::datatypes::DataType::Boolean, ArrowFillValue::Bool(fill_value)) => { + let bool_array = array.as_boolean(); + let filled_array: BooleanArray = bool_array + .iter() + .map(|v| v.filter(|&value| value != *fill_value)) + .collect(); + Arc::new(filled_array) + } + (arrow::datatypes::DataType::Int8, ArrowFillValue::Int8(fill_value)) => { + let int_array = array.as_primitive::(); + let filled_array: PrimitiveArray = int_array + .iter() + .map(|v| v.filter(|&value| value != *fill_value)) + .collect(); + Arc::new(filled_array) + } + (arrow::datatypes::DataType::Int16, ArrowFillValue::Int16(fill_value)) => { + let int_array = array.as_primitive::(); + let filled_array: PrimitiveArray = int_array + .iter() + .map(|v| v.filter(|&value| value != *fill_value)) + .collect(); + Arc::new(filled_array) + } + (arrow::datatypes::DataType::Int32, ArrowFillValue::Int32(fill_value)) => { + let int_array = array.as_primitive::(); + let filled_array: PrimitiveArray = int_array + .iter() + .map(|v| v.filter(|&value| value != *fill_value)) + .collect(); + Arc::new(filled_array) + } + (arrow::datatypes::DataType::Int64, ArrowFillValue::Int64(fill_value)) => { + let int_array = array.as_primitive::(); + let filled_array: PrimitiveArray = int_array + .iter() + .map(|v| v.filter(|&value| value != *fill_value)) + .collect(); + Arc::new(filled_array) + } + (arrow::datatypes::DataType::UInt8, ArrowFillValue::UInt8(fill_value)) => { + let int_array = array.as_primitive::(); + let filled_array: PrimitiveArray = int_array + .iter() + .map(|v| v.filter(|&value| value != *fill_value)) + .collect(); + Arc::new(filled_array) + } + (arrow::datatypes::DataType::UInt16, ArrowFillValue::UInt16(fill_value)) => { + let int_array = array.as_primitive::(); + let filled_array: PrimitiveArray = int_array + .iter() + .map(|v| v.filter(|&value| value != *fill_value)) + .collect(); + Arc::new(filled_array) + } + + (arrow::datatypes::DataType::UInt32, ArrowFillValue::UInt32(fill_value)) => { + let int_array = array.as_primitive::(); + let filled_array: PrimitiveArray = int_array + .iter() + .map(|v| v.filter(|&value| value != *fill_value)) + .collect(); + Arc::new(filled_array) + } + + (arrow::datatypes::DataType::UInt64, ArrowFillValue::UInt64(fill_value)) => { + let int_array = array.as_primitive::(); + let filled_array: PrimitiveArray = int_array + .iter() + .map(|v| v.filter(|&value| value != *fill_value)) + .collect(); + Arc::new(filled_array) + } + + (arrow::datatypes::DataType::Float32, ArrowFillValue::Float32(value)) => { + let float_array = array.as_primitive::(); + let filled_array: PrimitiveArray = float_array + .iter() + .map(|v| v.filter(|&v| v != *value)) + .collect(); + Arc::new(filled_array) + } + (arrow::datatypes::DataType::Float64, ArrowFillValue::Float64(value)) => { + let float_array = array.as_primitive::(); + let filled_array: PrimitiveArray = float_array + .iter() + .map(|v| v.filter(|&v| v != *value)) + .collect(); + Arc::new(filled_array) + } + (arrow::datatypes::DataType::Binary, ArrowFillValue::Bytes(fill_value)) => { + let binary_array = array.as_binary::(); + let filled_array: BinaryArray = binary_array + .iter() + .map(|v| v.filter(|&value| value != fill_value.as_slice())) + .collect(); + Arc::new(filled_array) + } + (arrow::datatypes::DataType::Utf8, ArrowFillValue::String(fill_value)) => { + let string_array = array.as_string::(); + let filled_array: StringArray = string_array + .iter() + .map(|v| v.filter(|&value| value != fill_value)) + .collect(); + + Arc::new(filled_array) + } + _ => { + return Err(format!( + "FillValueDecodingStep does not support data type: {:?}", + array.data_type() + )); + } + }; + + Ok(NdArrowArray::new(nullable_array, array.dimensions().clone()).unwrap()) + } +} + +pub struct CFDecoder { + // Array decoders to apply per array name + array_decoders: std::collections::HashMap>>, + decoded_schema: SchemaRef, +} + +impl Decoder for CFDecoder { + fn create( + group_reader: &crate::reader::AsyncArrowZarrGroupReader, + input_schema: SchemaRef, + ) -> Result { + let mut array_decoders: std::collections::HashMap< + String, + Vec>, + > = std::collections::HashMap::new(); + let mut fields = IndexMap::new(); + for field in input_schema.fields() { + fields.insert(field.name().clone(), field.as_ref().clone()); + } + + for (array_name, array_reader) in group_reader.arrays() { + // Get all the attributes that start with array_name + "." + let prefix = format!("{}.", array_name); + let array_attribute: Vec<_> = group_reader + .arrow_schema() + .fields() + .iter() + .filter(|field| field.name().starts_with(&prefix)) + .cloned() + .collect(); + + // Check for scale_factor and add_offset attributes (SCALE_FACTOR and ADD_OFFSET in CF conventions) + + let mut offset = None; + let mut scale_factor = None; + + for field in &array_attribute { + if *field.name() == format!("{}.scale_factor", array_name) { + match group_reader.attributes()[field.name()] { + crate::attributes::AttributeValue::Float64(value) => { + scale_factor = Some(value) + } + _ => { + return Err(format!( + "scale_factor attribute for array {} is not Float64", + array_name + )); + } + } + } else if *field.name() == format!("{}.add_offset", array_name) { + match group_reader.attributes()[field.name()] { + crate::attributes::AttributeValue::Float64(value) => offset = Some(value), + _ => { + return Err(format!( + "add_offset attribute for array {} is not Float64", + array_name + )); + } + } + } + } + + let decoding_pipeline = match (scale_factor, offset) { + (Some(scale), Some(offset)) => { + Some(Arc::new(ScaleFactorOffsetDecodingStep::new(scale, offset)) + as Arc) + } + (Some(scale), None) => { + Some(Arc::new(ScaleFactorOffsetDecodingStep::new(scale, 0.0)) + as Arc) + } + (None, Some(offset)) => { + Some(Arc::new(ScaleFactorOffsetDecodingStep::new(1.0, offset)) + as Arc) + } + _ => None, + }; + + if let Some(decoding_pipeline) = decoding_pipeline { + array_decoders + .entry(array_name.to_string()) + .or_default() + .push(decoding_pipeline.clone()); + // Update the field data type to Float64 + let array_field = fields.get_mut(array_name).unwrap(); + *array_field = array_field.clone().with_data_type( + decoding_pipeline.output_data_type(array_field.data_type().clone()), + ); + } + + // Check for cf time attributes (units) + for field in &array_attribute { + if *field.name() == format!("{}.units", array_name) + && let crate::attributes::AttributeValue::String(ref value) = + group_reader.attributes()[field.name()] + { + // Check if units is a valid CF time unit + if value.starts_with("seconds since") + || value.starts_with("minutes since") + || value.starts_with("hours since") + || value.starts_with("days since") + || value.starts_with("weeks since") + || value.starts_with("months since") + || value.starts_with("years since") + { + match TimeUnitsDecodingStep::new(value.clone()) { + Ok(decoder) => { + let decoder = Arc::new(decoder) as Arc; + array_decoders + .entry(array_name.to_string()) + .or_default() + .push(decoder.clone()); + // Update the field data type to Int64 (Arrow timestamp in milliseconds) + let array_field = fields.get_mut(array_name).unwrap(); + *array_field = array_field.clone().with_data_type( + decoder.output_data_type(array_field.data_type().clone()), + ); + } + Err(e) => { + tracing::warn!( + "Could not create TimeUnitsDecoder for array {}: {}", + array_name, + e + ); + } + } + } + } + } + } + + Ok(Self { + array_decoders, + decoded_schema: Arc::new(arrow::datatypes::Schema::new( + fields.into_values().collect::>(), + )), + }) + } + + fn decoded_schema(&self) -> SchemaRef { + self.decoded_schema.clone() + } + + fn decoding_array_pipeline( + &self, + array_name: &str, + ) -> Option>> { + self.array_decoders.get(array_name).cloned() + } +} + +pub struct ScaleFactorOffsetDecodingStep { + scale: f64, + offset: f64, +} + +impl ScaleFactorOffsetDecodingStep { + pub fn new(scale: f64, offset: f64) -> Self { + ScaleFactorOffsetDecodingStep { scale, offset } + } +} + +impl DecodingPipelineStep for ScaleFactorOffsetDecodingStep { + fn decode(&self, array: NdArrowArray) -> Result { + let arrow_array = array.as_arrow_array(); + + // Cast to float64 for processing + let float_array = + arrow::compute::cast(&arrow_array, &arrow::datatypes::DataType::Float64).unwrap(); + + let float_array = float_array.as_primitive::(); + + let applied_array: PrimitiveArray = + float_array.unary(|value| value * self.scale + self.offset); + + Ok(NdArrowArray::new( + Arc::new(applied_array) as arrow::array::ArrayRef, + array.dimensions().clone(), + ) + .unwrap()) + } + + fn output_data_type( + &self, + _input_data_type: arrow::datatypes::DataType, + ) -> arrow::datatypes::DataType { + arrow::datatypes::DataType::Float64 + } +} + +pub struct TimeUnitsDecodingStep { + unit: hifitime::Unit, + epoch: Epoch, +} + +impl TimeUnitsDecodingStep { + pub fn new(units_str: String) -> Result { + let unit = Self::extract_units(&units_str) + .ok_or(format!("Unsupported CF time units: {}", units_str))?; + let epoch = Self::extract_epoch(&units_str).ok_or(format!( + "Could not extract epoch from CF time units: {}", + units_str + ))?; + Ok(TimeUnitsDecodingStep { unit, epoch }) + } + + fn convert_array( + &self, + array: &PrimitiveArray, + unit: hifitime::Unit, + epoch: Epoch, + ) -> TimestampMillisecondArray { + let data: PrimitiveArray = + array.unary(|v| (epoch + (v * unit)).to_unix_milliseconds() as i64); + + // Reinterpret the Int64 array as TimestampMillisecond + data.reinterpret_cast::() + } + + fn extract_units(input: &str) -> Option { + let re = Regex::new(r"^(?P\w+) since").unwrap(); + re.captures(input) + .and_then(|caps| match caps["units"].to_string().as_str() { + "seconds" => Some(hifitime::Unit::Second), + "milliseconds" => Some(hifitime::Unit::Millisecond), + "microseconds" => Some(hifitime::Unit::Microsecond), + "nanoseconds" => Some(hifitime::Unit::Nanosecond), + "days" => Some(hifitime::Unit::Day), + "weeks" => Some(hifitime::Unit::Week), + _ => None, + }) + } + + /// Extracts the epoch date from a string like "days since -4713-11-24" + fn extract_epoch(input: &str) -> Option { + let re = Regex::new(r"since (?P-?\d{1,4}-\d{1,2}-\d{1,2})").unwrap(); + let result = re.captures(input).and_then(|caps| { + let epoch_str = caps["epoch"].to_string(); + let mut epoch = Epoch::from_str(&epoch_str).ok(); + + if epoch.is_none() && epoch_str == "-4713-01-01" { + epoch = Some(Epoch::from_jde_utc(0.0)); + } + + epoch + }); + + result + } +} + +impl DecodingPipelineStep for TimeUnitsDecodingStep { + fn output_data_type( + &self, + _input_data_type: arrow::datatypes::DataType, + ) -> arrow::datatypes::DataType { + arrow::datatypes::DataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None) + } + + fn decode(&self, array: NdArrowArray) -> Result { + let arrow_array = array.as_arrow_array(); + + // Cast to Int64 for processing + let int_array = + arrow::compute::cast(&arrow_array, &arrow::datatypes::DataType::Int64).unwrap(); + + let int_array = int_array.as_primitive::(); + + let applied_array = self.convert_array(int_array, self.unit, self.epoch); + + Ok(NdArrowArray::new( + Arc::new(applied_array) as arrow::array::ArrayRef, + array.dimensions().clone(), + ) + .unwrap()) + } +} diff --git a/beacon-arrow-zarr/src/lib.rs b/beacon-arrow-zarr/src/lib.rs new file mode 100644 index 00000000..12a62301 --- /dev/null +++ b/beacon-arrow-zarr/src/lib.rs @@ -0,0 +1,6 @@ +pub mod array_slice_pushdown; +pub mod attributes; +pub mod data_types; +pub mod decoder; +pub mod reader; +pub mod stream; diff --git a/beacon-arrow-zarr/src/reader.rs b/beacon-arrow-zarr/src/reader.rs new file mode 100644 index 00000000..465a2cc4 --- /dev/null +++ b/beacon-arrow-zarr/src/reader.rs @@ -0,0 +1,512 @@ +use std::{ + collections::{HashMap, HashSet}, + sync::Arc, +}; + +use arrow::array::BooleanArray; +use indexmap::IndexMap; +use nd_arrow_array::{ + NdArrowArray, + dimensions::{Dimension, Dimensions}, +}; +use zarrs::{array::Array, array_subset::ArraySubset, group::Group}; +use zarrs_storage::AsyncReadableListableStorageTraits; + +use crate::{ + array_slice_pushdown::ArraySlicePushDown, + attributes::AttributeValue, + decoder::{self, Decoder}, + stream::{ArrowZarrStreamComposer, ArrowZarrStreamComposerRef}, +}; + +pub struct AsyncArrowZarrGroupReader { + group: Arc>, + decoders: Vec>, + schema: arrow::datatypes::SchemaRef, + arrays: IndexMap>, + attributes: IndexMap, +} + +impl AsyncArrowZarrGroupReader { + pub async fn new( + group: Arc>, + ) -> Result { + let mut fields = Vec::new(); + + let mut attributes = IndexMap::new(); + let mut arrays = IndexMap::new(); + for (attr_name, attr_value) in group.attributes() { + if let Some(attr_value) = AttributeValue::from_json_value(attr_value) { + let arrow_data_type = attr_value.arrow_data_type(); + fields.push(arrow::datatypes::Field::new( + format!(".{}", attr_name), + arrow_data_type, + true, + )); + attributes.insert(format!(".{}", attr_name), attr_value); + } + } + + let group_path = group.path(); + + for array in group.async_child_arrays().await.unwrap() { + let array_json_attributes = array.attributes(); + let array_node_path = array.path().clone(); + let array_name = array_node_path + .as_str() + .strip_prefix(group_path.as_str()) + .unwrap() + .to_string(); + + // Strip the / from the start of the array name if present + let array_name = array_name.strip_prefix('/').unwrap_or(&array_name); + + for (attr_name, attr_value) in array_json_attributes { + if let Some(attr_value) = AttributeValue::from_json_value(attr_value) { + fields.push(arrow::datatypes::Field::new( + format!("{}.{}", array_name, attr_name), + attr_value.arrow_data_type(), + true, + )); + attributes.insert(format!("{}.{}", array_name, attr_name), attr_value); + } + } + + let zarr_data_type = array.data_type(); + if let Ok(arrow_data_type) = + crate::data_types::try_zarrs_dtype_to_arrow(zarr_data_type.clone()) + { + fields.push(arrow::datatypes::Field::new( + array_name, + arrow_data_type, + true, + )); + arrays.insert(array_name.to_string(), array); + } else { + // Unsupported data type, skip adding field + tracing::warn!( + "Skipping array {} with unsupported data type {:?}", + array_name, + zarr_data_type + ); + } + } + + let input_schema = Arc::new(arrow::datatypes::Schema::new(fields.clone())); + let mut self_reader = Self { + schema: input_schema.clone(), + group, + arrays, + attributes, + decoders: Vec::new(), + }; + + let mut decoders = Vec::new(); + // Create decoder for fill values + let fill_value_decoder = decoder::FillValueDecoder::create(&self_reader, input_schema)?; + self_reader.schema = fill_value_decoder.decoded_schema(); + decoders.push(Arc::new(fill_value_decoder) as Arc); + // Create decoder for CF conventions + let cf_decoder = decoder::CFDecoder::create(&self_reader, self_reader.schema.clone())?; + self_reader.schema = cf_decoder.decoded_schema(); + decoders.push(Arc::new(cf_decoder) as Arc); + self_reader.decoders = decoders; + Ok(self_reader) + } + + pub fn arrow_schema(&self) -> arrow::datatypes::SchemaRef { + self.schema.clone() + } + + pub fn arrays(&self) -> &IndexMap> { + &self.arrays + } + + pub fn attributes(&self) -> &IndexMap { + &self.attributes + } + + pub fn into_parallel_stream_composer>( + self, + projection: Option

, + array_slice_pushdowns: Option>, + ) -> Result { + let projected_schema = match projection { + Some(p) => Arc::new(self.schema.project(p.as_ref()).map_err(|e| e.to_string())?), + None => self.schema.clone(), + }; + + ArrowZarrStreamComposer::new(self, projected_schema, array_slice_pushdowns) + } + + pub fn read_attribute(&self, attribute_name: &str) -> Option { + self.attributes() + .get(attribute_name) + .map(|attr| attr.as_nd_arrow_array()) + } + + pub async fn read_array_full(&self, array_name: &str) -> Result, String> { + match self.arrays.get(array_name) { + Some(v) => { + let subset = v.subset_all(); + self.read_array(array_name, &subset).await + } + None => Ok(None), + } + } + + pub async fn read_array( + &self, + array_name: &str, + subset: &ArraySubset, + ) -> Result, String> { + let array_reader = if let Some(array) = self.arrays().get(array_name) { + array + } else { + return Ok(None); + }; + + let dimensions: Vec = array_reader + .dimension_names() + .as_ref() + .unwrap() + .iter() + .zip(subset.shape().iter()) + .map(|(dim, &len)| Dimension { + name: dim.clone().unwrap(), + size: len as usize, + }) + .collect::>(); + let dimensions = Dimensions::new(dimensions); + + let mut maybe_nd_array = match array_reader.data_type() { + zarrs::array::DataType::Bool => { + let array = array_reader + .async_retrieve_array_subset_ndarray::(subset) + .await + .map_err(|e| e.to_string())?; + + let arrow_boolean_array = BooleanArray::from(array.into_raw_vec_and_offset().0); + + let nd_array = NdArrowArray::new(Arc::new(arrow_boolean_array), dimensions) + .map_err(|e| e.to_string())?; + Ok(Some(nd_array)) + } + zarrs::array::DataType::Int8 => { + let array = array_reader + .async_retrieve_array_subset_ndarray::(subset) + .await + .map_err(|e| e.to_string())?; + + let arrow_int8_array = + arrow::array::Int8Array::from(array.into_raw_vec_and_offset().0); + + let nd_array = NdArrowArray::new(Arc::new(arrow_int8_array), dimensions) + .map_err(|e| e.to_string())?; + Ok(Some(nd_array)) + } + zarrs::array::DataType::Int16 => { + let array = array_reader + .async_retrieve_array_subset_ndarray::(subset) + .await + .map_err(|e| e.to_string())?; + + let arrow_int16_array = + arrow::array::Int16Array::from(array.into_raw_vec_and_offset().0); + + let nd_array = NdArrowArray::new(Arc::new(arrow_int16_array), dimensions) + .map_err(|e| e.to_string())?; + Ok(Some(nd_array)) + } + zarrs::array::DataType::Int32 => { + let array = array_reader + .async_retrieve_array_subset_ndarray::(subset) + .await + .map_err(|e| e.to_string())?; + + let arrow_int32_array = + arrow::array::Int32Array::from(array.into_raw_vec_and_offset().0); + + let nd_array = NdArrowArray::new(Arc::new(arrow_int32_array), dimensions) + .map_err(|e| e.to_string())?; + Ok(Some(nd_array)) + } + zarrs::array::DataType::Int64 => { + let array = array_reader + .async_retrieve_array_subset_ndarray::(subset) + .await + .map_err(|e| e.to_string())?; + + let arrow_int64_array = + arrow::array::Int64Array::from(array.into_raw_vec_and_offset().0); + + let nd_array = NdArrowArray::new(Arc::new(arrow_int64_array), dimensions) + .map_err(|e| e.to_string())?; + Ok(Some(nd_array)) + } + zarrs::array::DataType::UInt8 => { + let array = array_reader + .async_retrieve_array_subset_ndarray::(subset) + .await + .map_err(|e| e.to_string())?; + + let arrow_uint8_array = + arrow::array::UInt8Array::from(array.into_raw_vec_and_offset().0); + + let nd_array = NdArrowArray::new(Arc::new(arrow_uint8_array), dimensions) + .map_err(|e| e.to_string())?; + Ok(Some(nd_array)) + } + zarrs::array::DataType::UInt16 => { + let array = array_reader + .async_retrieve_array_subset_ndarray::(subset) + .await + .map_err(|e| e.to_string())?; + + let arrow_uint16_array = + arrow::array::UInt16Array::from(array.into_raw_vec_and_offset().0); + + let nd_array = NdArrowArray::new(Arc::new(arrow_uint16_array), dimensions) + .map_err(|e| e.to_string())?; + Ok(Some(nd_array)) + } + zarrs::array::DataType::UInt32 => { + let array = array_reader + .async_retrieve_array_subset_ndarray::(subset) + .await + .map_err(|e| e.to_string())?; + + let arrow_uint32_array = + arrow::array::UInt32Array::from(array.into_raw_vec_and_offset().0); + + let nd_array = NdArrowArray::new(Arc::new(arrow_uint32_array), dimensions) + .map_err(|e| e.to_string())?; + Ok(Some(nd_array)) + } + zarrs::array::DataType::UInt64 => { + let array = array_reader + .async_retrieve_array_subset_ndarray::(subset) + .await + .map_err(|e| e.to_string())?; + + let arrow_uint64_array = + arrow::array::UInt64Array::from(array.into_raw_vec_and_offset().0); + + let nd_array = NdArrowArray::new(Arc::new(arrow_uint64_array), dimensions) + .map_err(|e| e.to_string())?; + Ok(Some(nd_array)) + } + zarrs::array::DataType::Float32 => { + let array = array_reader + .async_retrieve_array_subset_ndarray::(subset) + .await + .map_err(|e| e.to_string())?; + + let arrow_float32_array = + arrow::array::Float32Array::from(array.into_raw_vec_and_offset().0); + + let nd_array = NdArrowArray::new(Arc::new(arrow_float32_array), dimensions) + .map_err(|e| e.to_string())?; + Ok(Some(nd_array)) + } + zarrs::array::DataType::Float64 => { + let array = array_reader + .async_retrieve_array_subset_ndarray::(subset) + .await + .map_err(|e| e.to_string())?; + + let arrow_float64_array = + arrow::array::Float64Array::from(array.into_raw_vec_and_offset().0); + + let nd_array = NdArrowArray::new(Arc::new(arrow_float64_array), dimensions) + .map_err(|e| e.to_string())?; + Ok(Some(nd_array)) + } + zarrs::array::DataType::String => { + let array = array_reader + .async_retrieve_array_subset_ndarray::(subset) + .await + .map_err(|e| e.to_string())?; + + let arrow_string_array = + arrow::array::StringArray::from(array.into_raw_vec_and_offset().0); + + let nd_array = NdArrowArray::new(Arc::new(arrow_string_array), dimensions) + .map_err(|e| e.to_string())?; + Ok(Some(nd_array)) + } + zarrs::array::DataType::Bytes => { + let array = array_reader + .async_retrieve_array_subset_ndarray::>(subset) + .await + .map_err(|e| e.to_string())?; + + let arrow_binary_array = + arrow::array::BinaryArray::from_iter_values(array.into_raw_vec_and_offset().0); + + let nd_array = NdArrowArray::new(Arc::new(arrow_binary_array), dimensions) + .map_err(|e| e.to_string())?; + Ok(Some(nd_array)) + } + zarr_data_type => Err(format!("Unsupported Zarrs data type: {:?}", zarr_data_type)), + }; + + maybe_nd_array = if let Ok(Some(mut nd_array)) = maybe_nd_array { + // Apply decoders + for decoder in &self.decoders { + nd_array = decoder.decode_array(array_name, nd_array)?; + } + + Ok(Some(nd_array)) + } else { + maybe_nd_array + }; + + maybe_nd_array + } +} + +#[cfg(test)] +mod tests { + + use std::sync::Arc; + + use futures::StreamExt; + use object_store::local::LocalFileSystem; + use zarrs::{array::Array, array_subset::ArraySubset, group::Group}; + use zarrs_object_store::AsyncObjectStore; + use zarrs_storage::AsyncReadableListableStorage; + + #[tokio::test] + async fn test_name() { + let local_fs = LocalFileSystem::new_with_prefix("./test_files").unwrap(); + + let object_store = AsyncObjectStore::new(local_fs); + + let zarr_store: AsyncReadableListableStorage = Arc::new(object_store); + + let group = Group::async_open(zarr_store.clone(), "/gridded-example.zarr") + .await + .unwrap(); + + println!("{:?}", group.attributes()); + + println!("{:?}", group.async_child_array_paths().await.unwrap()); + println!("{:?}", group.async_child_group_paths().await.unwrap()); + + let sst_arr = Array::async_open(zarr_store.clone(), "/gridded-example.zarr/analysed_sst") + .await + .unwrap(); + println!( + "SST: {:?}", + sst_arr + .chunk_grid() + .iter_chunk_indices() + .collect::>() + ); + + println!("SST attrs: {:?}", sst_arr.attributes()); + println!("SST dims: {:?}", sst_arr.dimension_names()); + + let y = sst_arr + .chunk_grid() + .iter_chunk_indices() + .collect::>(); + let chunk_grid = sst_arr.chunk_grid().clone(); + + println!("SST chunk subset: {:?}", sst_arr.chunk_subset(&[0, 3, 1])); + println!("SST chunk subset: {:?}", chunk_grid.subset(&[0, 3, 1])); + + let lat_arr = Array::async_open(zarr_store.clone(), "/gridded-example.zarr/lat") + .await + .unwrap(); + println!("Lat: {:?}", lat_arr.chunk_grid()); + println!("Lat data type: {:?}", lat_arr.data_type()); + + let arr = lat_arr + .async_retrieve_array_subset_ndarray::(&ArraySubset::new_with_ranges(&[(0..10)])) + .await + .unwrap(); + + println!("Lat arr: {:?}", arr); + + let time_arr = Array::async_open(zarr_store.clone(), "/gridded-example.zarr/time") + .await + .unwrap(); + + println!("Time: {:?}", time_arr.chunk_grid()); + println!("Time shape: {:?}", time_arr.chunk_grid_shape()); + println!("Time attrs: {:?}", time_arr.attributes()); + println!("Time data type: {:?}", time_arr.data_type()); + } + + #[tokio::test] + async fn test_arrow_schema() { + let local_fs = LocalFileSystem::new_with_prefix("./test_files").unwrap(); + + let object_store = AsyncObjectStore::new(local_fs); + + let zarr_store: AsyncReadableListableStorage = Arc::new(object_store); + + let group = Group::async_open(zarr_store.clone(), "/gridded-example.zarr") + .await + .unwrap(); + + let reader = super::AsyncArrowZarrGroupReader::new(Arc::new(group)) + .await + .unwrap(); + + let schema = reader.arrow_schema(); + for field in schema.fields() { + println!("Field: {:?} : {:?}", field.name(), field.data_type()); + } + } + + #[tokio::test] + async fn test_read_stream() { + let local_fs = LocalFileSystem::new_with_prefix("./test_files").unwrap(); + + let object_store = AsyncObjectStore::new(local_fs); + + let zarr_store: AsyncReadableListableStorage = Arc::new(object_store); + + let group = Group::async_open(zarr_store.clone(), "/gridded-example.zarr") + .await + .unwrap(); + + let reader = super::AsyncArrowZarrGroupReader::new(Arc::new(group)) + .await + .unwrap(); + + let longitude_idx = reader + .schema + .index_of("lon") + .expect("longitude field not found in schema"); + let latitude_idx = reader + .schema + .index_of("lat") + .expect("latitude field not found in schema"); + + let analysed_sst_idx = reader + .schema + .index_of("analysed_sst") + .expect("analysed_sst field not found in schema"); + + let mut stream = reader + .into_parallel_stream_composer( + Some(&[longitude_idx, latitude_idx, analysed_sst_idx]), + None, + ) + .unwrap() + .pollable_shared_stream(); + + let next_batch = stream.next().await; + + arrow::util::pretty::print_batches(&[next_batch + .unwrap() + .unwrap() + .to_arrow_record_batch() + .unwrap()]) + .unwrap(); + } +} diff --git a/beacon-arrow-zarr/src/stream.rs b/beacon-arrow-zarr/src/stream.rs new file mode 100644 index 00000000..4010be3a --- /dev/null +++ b/beacon-arrow-zarr/src/stream.rs @@ -0,0 +1,256 @@ +use std::{collections::HashMap, pin::Pin, sync::Arc, task::Poll}; + +use futures::{Future, Stream}; +use indexmap::IndexMap; +use nd_arrow_array::{NdArrowArray, batch::NdRecordBatch}; +use zarrs::{array::ChunkGrid, array_subset::ArraySubset}; + +use crate::{array_slice_pushdown::ArraySlicePushDown, reader::AsyncArrowZarrGroupReader}; + +pub type ArrowZarrStreamComposerRef = Arc; + +pub struct ArrowZarrStreamComposer { + group_reader: crate::reader::AsyncArrowZarrGroupReader, + projected_schema: arrow::datatypes::SchemaRef, + // Streaming State + chunk_grid: Option, + dimension_names: Vec, + dimension_sizes: Vec, + readable_chunks: crossbeam::queue::SegQueue, + array_slice_pushdowns: Option>, +} + +impl ArrowZarrStreamComposer { + pub fn new( + group_reader: AsyncArrowZarrGroupReader, + projected_schema: arrow::datatypes::SchemaRef, + array_slice_pushdowns: Option>, + ) -> Result { + // Get all the variables + let variables = projected_schema + .fields() + .iter() + .filter_map(|f| { + if !f.name().contains('.') { + Some(f.name().clone()) + } else { + None + } + }) + .collect::>(); + + // Find the variable with the most dimensions + let max_dims_var = variables + .iter() + .max_by_key(|var_name| { + group_reader + .arrays() + .get(var_name.as_str()) + .map(|array| array.shape().len()) + .unwrap_or(0) + }) + .cloned(); + + let mut chunk_grid = None; + let mut dimension_names = vec![]; + let mut dimension_sizes = vec![]; + let queue = crossbeam::queue::SegQueue::new(); + + if let Some(var_name) = max_dims_var { + let var = group_reader.arrays().get(var_name.as_str()).unwrap(); + let chunk_shape = var.chunk_grid(); + let var_dimension_names = var + .dimension_names() + .as_ref() + .unwrap() + .iter() + .map(|d| d.clone().unwrap()) + .collect::>(); + + chunk_grid = Some(chunk_shape.clone()); + dimension_names = var_dimension_names; + dimension_sizes = var.shape().to_vec(); + } + + if let Some(chunk_grid) = &chunk_grid { + // Generate all chunk indices + for chunk_indices in chunk_grid.iter_chunk_indices() { + queue.push(chunk_indices); + } + } else { + // We should at least have one chunk to read, we can add an empty as it should only contain scalar attributes and arrays + queue.push(Vec::new()); + } + + Ok(Arc::new(Self { + group_reader, + projected_schema, + chunk_grid, + dimension_names, + dimension_sizes, + readable_chunks: queue, + array_slice_pushdowns, + })) + } + + pub fn generate_array_subset(&self, chunk_indices: &ChunkIndices) -> ArraySubset { + match self.chunk_grid.as_ref() { + Some(chunk_grid) => chunk_grid.subset(chunk_indices).unwrap().unwrap(), + None => ArraySubset::new_empty(0), // Scalar array + } + } + + pub fn pollable_shared_stream(self: &Arc) -> ArrowZarrStream { + ArrowZarrStream { + composer: self.clone(), + ongoing: None, + } + } +} + +pub type ChunkIndices = Vec; + +type ArrowZarrStreamFuture = + Pin>> + Send>>; + +#[pin_project::pin_project] +pub struct ArrowZarrStream { + composer: Arc, + // Pinned state + #[pin] + ongoing: Option, +} + +impl ArrowZarrStream { + pub async fn next_chunk( + composer: Arc, + ) -> Option> { + let chunk_indices = match composer.readable_chunks.pop() { + Some(chunk_indices) => chunk_indices, + None => return None, // Queue is empty + }; + let array_subset = composer.generate_array_subset(&chunk_indices); + + // Blend dimension_names and array_subset ranges into a map (this will wrap ranges that exceed dimension sizes) + let mut dimension_subset_ranges: IndexMap> = composer + .dimension_names + .iter() + .zip(array_subset.to_ranges()) + .zip(composer.dimension_sizes.iter()) + .map(|((name, range), dim_size)| { + if range.end > *dim_size { + (name.clone(), range.start..*dim_size) + } else { + (name.clone(), range) + } + }) + .collect(); + if let Some(pushdowns) = &composer.array_slice_pushdowns { + // Apply array slice pushdowns + for (dim_name, pushdown) in pushdowns.iter() { + if let Some(range) = dimension_subset_ranges.get_mut(dim_name) { + if let Some(overlap) = + pushdown.overlapping_range(range.start as usize..range.end as usize) + { + *range = overlap.start as u64..overlap.end as u64; + } else { + // No overlap, return empty batch + let mut arrays = vec![]; + let mut fields = vec![]; + for field in composer.projected_schema.fields() { + fields.push(field.as_ref().clone()); + arrays.push(NdArrowArray::new_null_scalar(Some( + field.data_type().clone(), + ))); + } + + let empty_batch = + NdRecordBatch::new(fields, arrays).map_err(|e| e.to_string()); + return Some(empty_batch); + } + } + } + } + tracing::debug!( + "Reading chunk {:?} with subset ranges: {:?}", + chunk_indices, + dimension_subset_ranges + ); + let mut arrays = vec![]; + let mut fields = vec![]; + for column in composer.projected_schema.fields().iter() { + let array_name = column.name(); + fields.push(column.as_ref().clone()); + + if array_name.contains('.') { + // Reading attribute + if let Some(attribute) = composer.group_reader.read_attribute(array_name) { + arrays.push(attribute); + } else { + return Some(Err(format!( + "Attribute {} not found in Zarr group.", + array_name + ))); + } + } else { + // Reading array + let subset = if let Some(dimension_names) = + composer.group_reader.arrays()[array_name.as_str()].dimension_names() + { + let ranges = dimension_names + .iter() + .map(|dim_name| { + dimension_subset_ranges + .get(dim_name.as_ref().unwrap().as_str()) + .unwrap() + .clone() + }) + .collect::>(); + + ArraySubset::new_with_ranges(&ranges) + } else { + // Scalar array + ArraySubset::new_empty(0) + }; + let array = composer + .group_reader + .read_array(array_name, &subset) + .await + .unwrap() + .unwrap(); + + arrays.push(array); + } + } + + let record_batch = NdRecordBatch::new(fields, arrays).map_err(|e| e.to_string()); + + Some(record_batch) + } +} + +impl Stream for ArrowZarrStream { + type Item = Result; + + fn poll_next( + self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + let mut pinned_stream = self.project(); + let composer = pinned_stream.composer.clone(); + + if pinned_stream.ongoing.is_none() { + *pinned_stream.ongoing = Some(Box::pin(Self::next_chunk(composer)) as _); + } + + let fut = pinned_stream.ongoing.as_mut().as_pin_mut().unwrap(); + match fut.poll(cx) { + Poll::Pending => Poll::Pending, + Poll::Ready(Some(value)) => { + *pinned_stream.ongoing = None; + Poll::Ready(Some(value)) + } + Poll::Ready(None) => Poll::Ready(None), + } + } +} diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/analysed_sst/c/0/0/0 b/beacon-arrow-zarr/test_files/gridded-example.zarr/analysed_sst/c/0/0/0 new file mode 100644 index 00000000..fb949fc5 Binary files /dev/null and b/beacon-arrow-zarr/test_files/gridded-example.zarr/analysed_sst/c/0/0/0 differ diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/analysed_sst/c/0/0/1 b/beacon-arrow-zarr/test_files/gridded-example.zarr/analysed_sst/c/0/0/1 new file mode 100644 index 00000000..b954f909 Binary files /dev/null and b/beacon-arrow-zarr/test_files/gridded-example.zarr/analysed_sst/c/0/0/1 differ diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/analysed_sst/c/0/1/0 b/beacon-arrow-zarr/test_files/gridded-example.zarr/analysed_sst/c/0/1/0 new file mode 100644 index 00000000..965a478b Binary files /dev/null and b/beacon-arrow-zarr/test_files/gridded-example.zarr/analysed_sst/c/0/1/0 differ diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/analysed_sst/c/0/1/1 b/beacon-arrow-zarr/test_files/gridded-example.zarr/analysed_sst/c/0/1/1 new file mode 100644 index 00000000..92981f4f Binary files /dev/null and b/beacon-arrow-zarr/test_files/gridded-example.zarr/analysed_sst/c/0/1/1 differ diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/analysed_sst/c/0/2/0 b/beacon-arrow-zarr/test_files/gridded-example.zarr/analysed_sst/c/0/2/0 new file mode 100644 index 00000000..fbe7f5cc Binary files /dev/null and b/beacon-arrow-zarr/test_files/gridded-example.zarr/analysed_sst/c/0/2/0 differ diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/analysed_sst/c/0/2/1 b/beacon-arrow-zarr/test_files/gridded-example.zarr/analysed_sst/c/0/2/1 new file mode 100644 index 00000000..84b039c6 Binary files /dev/null and b/beacon-arrow-zarr/test_files/gridded-example.zarr/analysed_sst/c/0/2/1 differ diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/analysed_sst/c/0/3/0 b/beacon-arrow-zarr/test_files/gridded-example.zarr/analysed_sst/c/0/3/0 new file mode 100644 index 00000000..4a2f1ad4 Binary files /dev/null and b/beacon-arrow-zarr/test_files/gridded-example.zarr/analysed_sst/c/0/3/0 differ diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/analysed_sst/c/0/3/1 b/beacon-arrow-zarr/test_files/gridded-example.zarr/analysed_sst/c/0/3/1 new file mode 100644 index 00000000..5e829e45 Binary files /dev/null and b/beacon-arrow-zarr/test_files/gridded-example.zarr/analysed_sst/c/0/3/1 differ diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/analysed_sst/zarr.json b/beacon-arrow-zarr/test_files/gridded-example.zarr/analysed_sst/zarr.json new file mode 100644 index 00000000..69af4005 --- /dev/null +++ b/beacon-arrow-zarr/test_files/gridded-example.zarr/analysed_sst/zarr.json @@ -0,0 +1,61 @@ +{ + "shape": [ + 1, + 1208, + 1920 + ], + "data_type": "int16", + "chunk_grid": { + "name": "regular", + "configuration": { + "chunk_shape": [ + 1, + 302, + 960 + ] + } + }, + "chunk_key_encoding": { + "name": "default", + "configuration": { + "separator": "/" + } + }, + "fill_value": 0, + "codecs": [ + { + "name": "bytes", + "configuration": { + "endian": "little" + } + }, + { + "name": "zstd", + "configuration": { + "level": 0, + "checksum": false + } + } + ], + "attributes": { + "long_name": "analysed sea surface temperature", + "standard_name": "sea_surface_foundation_temperature", + "type": "foundation", + "units": "kelvin", + "valid_min": -300, + "valid_max": 4500, + "source": "EUR-L2P-ATS_NR_2P,UPA-L2P-ATS_NR_2P,JPL-L2P-MODIS_A,JPL-L2P-MODIS_T,EUR-L2P-NAR17_SST,EUR-L2P-NAR18_SST,EUR-L2P-SEVIRI_SST,EUR-L2P-AVHRR_METOP_A,EUR-L2P-AVHRR_METOP_B,NAVO-L2P-VIIRS_NPP,EUR-L2P-SEVIRI_SST_Meteosat11,MAR-L2P-SLSTR_S3A,MAR-L2P-SLSTR_S3B,STAR-L2P-VIIRS_N20-ACSPO_V2.80", + "comment": "Optimal interpolation (OI) at Ultrahigh resolution of sst measurements from infrared sensors", + "add_offset": 273.1499938964844, + "scale_factor": 0.009999999776482582, + "_FillValue": -32768 + }, + "dimension_names": [ + "time", + "lat", + "lon" + ], + "zarr_format": 3, + "node_type": "array", + "storage_transformers": [] +} \ No newline at end of file diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/analysis_error/c/0/0/0 b/beacon-arrow-zarr/test_files/gridded-example.zarr/analysis_error/c/0/0/0 new file mode 100644 index 00000000..32e07e94 Binary files /dev/null and b/beacon-arrow-zarr/test_files/gridded-example.zarr/analysis_error/c/0/0/0 differ diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/analysis_error/c/0/0/1 b/beacon-arrow-zarr/test_files/gridded-example.zarr/analysis_error/c/0/0/1 new file mode 100644 index 00000000..a9a32956 Binary files /dev/null and b/beacon-arrow-zarr/test_files/gridded-example.zarr/analysis_error/c/0/0/1 differ diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/analysis_error/c/0/1/0 b/beacon-arrow-zarr/test_files/gridded-example.zarr/analysis_error/c/0/1/0 new file mode 100644 index 00000000..cadece7f Binary files /dev/null and b/beacon-arrow-zarr/test_files/gridded-example.zarr/analysis_error/c/0/1/0 differ diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/analysis_error/c/0/1/1 b/beacon-arrow-zarr/test_files/gridded-example.zarr/analysis_error/c/0/1/1 new file mode 100644 index 00000000..980c9a69 Binary files /dev/null and b/beacon-arrow-zarr/test_files/gridded-example.zarr/analysis_error/c/0/1/1 differ diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/analysis_error/c/0/2/0 b/beacon-arrow-zarr/test_files/gridded-example.zarr/analysis_error/c/0/2/0 new file mode 100644 index 00000000..a2568e92 Binary files /dev/null and b/beacon-arrow-zarr/test_files/gridded-example.zarr/analysis_error/c/0/2/0 differ diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/analysis_error/c/0/2/1 b/beacon-arrow-zarr/test_files/gridded-example.zarr/analysis_error/c/0/2/1 new file mode 100644 index 00000000..dc7b4e3e Binary files /dev/null and b/beacon-arrow-zarr/test_files/gridded-example.zarr/analysis_error/c/0/2/1 differ diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/analysis_error/c/0/3/0 b/beacon-arrow-zarr/test_files/gridded-example.zarr/analysis_error/c/0/3/0 new file mode 100644 index 00000000..61d3ec27 Binary files /dev/null and b/beacon-arrow-zarr/test_files/gridded-example.zarr/analysis_error/c/0/3/0 differ diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/analysis_error/c/0/3/1 b/beacon-arrow-zarr/test_files/gridded-example.zarr/analysis_error/c/0/3/1 new file mode 100644 index 00000000..10ff4e67 Binary files /dev/null and b/beacon-arrow-zarr/test_files/gridded-example.zarr/analysis_error/c/0/3/1 differ diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/analysis_error/zarr.json b/beacon-arrow-zarr/test_files/gridded-example.zarr/analysis_error/zarr.json new file mode 100644 index 00000000..f21e202d --- /dev/null +++ b/beacon-arrow-zarr/test_files/gridded-example.zarr/analysis_error/zarr.json @@ -0,0 +1,59 @@ +{ + "shape": [ + 1, + 1208, + 1920 + ], + "data_type": "int16", + "chunk_grid": { + "name": "regular", + "configuration": { + "chunk_shape": [ + 1, + 302, + 960 + ] + } + }, + "chunk_key_encoding": { + "name": "default", + "configuration": { + "separator": "/" + } + }, + "fill_value": 0, + "codecs": [ + { + "name": "bytes", + "configuration": { + "endian": "little" + } + }, + { + "name": "zstd", + "configuration": { + "level": 0, + "checksum": false + } + } + ], + "attributes": { + "long_name": "estimated error standard deviation of analysed_sst", + "units": "kelvin", + "valid_min": 0, + "valid_max": 10000, + "source": "EUR-L2P-ATS_NR_2P,UPA-L2P-ATS_NR_2P,JPL-L2P-MODIS_A,JPL-L2P-MODIS_T,EUR-L2P-NAR17_SST,EUR-L2P-NAR18_SST,EUR-L2P-SEVIRI_SST,EUR-L2P-AVHRR_METOP_A,EUR-L2P-AVHRR_METOP_B,NAVO-L2P-VIIRS_NPP,EUR-L2P-SEVIRI_SST_Meteosat11,MAR-L2P-SLSTR_S3A,MAR-L2P-SLSTR_S3B,STAR-L2P-VIIRS_N20-ACSPO_V2.80", + "comment": "Error obtained by Optimal interpolation (OI) at Ultrahigh resolution of sst measurements from infrared sensors", + "add_offset": 0.0, + "scale_factor": 0.009999999776482582, + "_FillValue": -32768 + }, + "dimension_names": [ + "time", + "lat", + "lon" + ], + "zarr_format": 3, + "node_type": "array", + "storage_transformers": [] +} \ No newline at end of file diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/lat/c/0 b/beacon-arrow-zarr/test_files/gridded-example.zarr/lat/c/0 new file mode 100644 index 00000000..fb6d5039 Binary files /dev/null and b/beacon-arrow-zarr/test_files/gridded-example.zarr/lat/c/0 differ diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/lat/zarr.json b/beacon-arrow-zarr/test_files/gridded-example.zarr/lat/zarr.json new file mode 100644 index 00000000..54fd729f --- /dev/null +++ b/beacon-arrow-zarr/test_files/gridded-example.zarr/lat/zarr.json @@ -0,0 +1,52 @@ +{ + "shape": [ + 1208 + ], + "data_type": "float32", + "chunk_grid": { + "name": "regular", + "configuration": { + "chunk_shape": [ + 1208 + ] + } + }, + "chunk_key_encoding": { + "name": "default", + "configuration": { + "separator": "/" + } + }, + "fill_value": 0.0, + "codecs": [ + { + "name": "bytes", + "configuration": { + "endian": "little" + } + }, + { + "name": "zstd", + "configuration": { + "level": 0, + "checksum": false + } + } + ], + "attributes": { + "long_name": "latitude", + "standard_name": "latitude", + "axis": "Y", + "units": "degrees_north", + "comment": "geographical coordinates, WGS84 projection", + "valid_min": 37.0, + "valid_max": 49.0, + "_FillValue": "AAAAAAAA+H8=" + }, + "dimension_names": [ + "lat" + ], + "zarr_format": 3, + "node_type": "array", + "storage_transformers": [] +} \ No newline at end of file diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/lon/c/0 b/beacon-arrow-zarr/test_files/gridded-example.zarr/lon/c/0 new file mode 100644 index 00000000..322cddab Binary files /dev/null and b/beacon-arrow-zarr/test_files/gridded-example.zarr/lon/c/0 differ diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/lon/zarr.json b/beacon-arrow-zarr/test_files/gridded-example.zarr/lon/zarr.json new file mode 100644 index 00000000..ffd46435 --- /dev/null +++ b/beacon-arrow-zarr/test_files/gridded-example.zarr/lon/zarr.json @@ -0,0 +1,52 @@ +{ + "shape": [ + 1920 + ], + "data_type": "float32", + "chunk_grid": { + "name": "regular", + "configuration": { + "chunk_shape": [ + 1920 + ] + } + }, + "chunk_key_encoding": { + "name": "default", + "configuration": { + "separator": "/" + } + }, + "fill_value": 0.0, + "codecs": [ + { + "name": "bytes", + "configuration": { + "endian": "little" + } + }, + { + "name": "zstd", + "configuration": { + "level": 0, + "checksum": false + } + } + ], + "attributes": { + "long_name": "longitude", + "standard_name": "longitude", + "axis": "X", + "units": "degrees_east", + "comment": "geographical coordinates, WGS84 projection", + "valid_min": 25.0, + "valid_max": 43.0, + "_FillValue": "AAAAAAAA+H8=" + }, + "dimension_names": [ + "lon" + ], + "zarr_format": 3, + "node_type": "array", + "storage_transformers": [] +} \ No newline at end of file diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/mask/c/0/0/0 b/beacon-arrow-zarr/test_files/gridded-example.zarr/mask/c/0/0/0 new file mode 100644 index 00000000..17a39e13 Binary files /dev/null and b/beacon-arrow-zarr/test_files/gridded-example.zarr/mask/c/0/0/0 differ diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/mask/c/0/0/1 b/beacon-arrow-zarr/test_files/gridded-example.zarr/mask/c/0/0/1 new file mode 100644 index 00000000..76fd11e5 Binary files /dev/null and b/beacon-arrow-zarr/test_files/gridded-example.zarr/mask/c/0/0/1 differ diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/mask/c/0/1/0 b/beacon-arrow-zarr/test_files/gridded-example.zarr/mask/c/0/1/0 new file mode 100644 index 00000000..190eed43 Binary files /dev/null and b/beacon-arrow-zarr/test_files/gridded-example.zarr/mask/c/0/1/0 differ diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/mask/c/0/1/1 b/beacon-arrow-zarr/test_files/gridded-example.zarr/mask/c/0/1/1 new file mode 100644 index 00000000..efe5b467 Binary files /dev/null and b/beacon-arrow-zarr/test_files/gridded-example.zarr/mask/c/0/1/1 differ diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/mask/c/0/2/0 b/beacon-arrow-zarr/test_files/gridded-example.zarr/mask/c/0/2/0 new file mode 100644 index 00000000..e2f6eaac Binary files /dev/null and b/beacon-arrow-zarr/test_files/gridded-example.zarr/mask/c/0/2/0 differ diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/mask/c/0/2/1 b/beacon-arrow-zarr/test_files/gridded-example.zarr/mask/c/0/2/1 new file mode 100644 index 00000000..4a6e47e4 Binary files /dev/null and b/beacon-arrow-zarr/test_files/gridded-example.zarr/mask/c/0/2/1 differ diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/mask/c/0/3/0 b/beacon-arrow-zarr/test_files/gridded-example.zarr/mask/c/0/3/0 new file mode 100644 index 00000000..a94bb8d3 Binary files /dev/null and b/beacon-arrow-zarr/test_files/gridded-example.zarr/mask/c/0/3/0 differ diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/mask/c/0/3/1 b/beacon-arrow-zarr/test_files/gridded-example.zarr/mask/c/0/3/1 new file mode 100644 index 00000000..ddcbbcae Binary files /dev/null and b/beacon-arrow-zarr/test_files/gridded-example.zarr/mask/c/0/3/1 differ diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/mask/zarr.json b/beacon-arrow-zarr/test_files/gridded-example.zarr/mask/zarr.json new file mode 100644 index 00000000..908e44a2 --- /dev/null +++ b/beacon-arrow-zarr/test_files/gridded-example.zarr/mask/zarr.json @@ -0,0 +1,60 @@ +{ + "shape": [ + 1, + 1208, + 1920 + ], + "data_type": "int8", + "chunk_grid": { + "name": "regular", + "configuration": { + "chunk_shape": [ + 1, + 302, + 960 + ] + } + }, + "chunk_key_encoding": { + "name": "default", + "configuration": { + "separator": "/" + } + }, + "fill_value": 0, + "codecs": [ + { + "name": "bytes" + }, + { + "name": "zstd", + "configuration": { + "level": 0, + "checksum": false + } + } + ], + "attributes": { + "long_name": "sea/land/lake/ice field composite mask", + "flag_values": [ + 1, + 2, + 3, + 4 + ], + "flag_meanings": "sea land lake ice", + "valid_min": 1, + "valid_max": 4, + "source_mask": "----", + "comment": "b0:1= sea; b1:1= land; b2:1= lake; b3:1= sea ice; b4-7: reserved for grid mask", + "_FillValue": -128 + }, + "dimension_names": [ + "time", + "lat", + "lon" + ], + "zarr_format": 3, + "node_type": "array", + "storage_transformers": [] +} \ No newline at end of file diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/sea_ice_fraction/c/0/0/0 b/beacon-arrow-zarr/test_files/gridded-example.zarr/sea_ice_fraction/c/0/0/0 new file mode 100644 index 00000000..5716f111 Binary files /dev/null and b/beacon-arrow-zarr/test_files/gridded-example.zarr/sea_ice_fraction/c/0/0/0 differ diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/sea_ice_fraction/c/0/0/1 b/beacon-arrow-zarr/test_files/gridded-example.zarr/sea_ice_fraction/c/0/0/1 new file mode 100644 index 00000000..5716f111 Binary files /dev/null and b/beacon-arrow-zarr/test_files/gridded-example.zarr/sea_ice_fraction/c/0/0/1 differ diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/sea_ice_fraction/c/0/1/0 b/beacon-arrow-zarr/test_files/gridded-example.zarr/sea_ice_fraction/c/0/1/0 new file mode 100644 index 00000000..5716f111 Binary files /dev/null and b/beacon-arrow-zarr/test_files/gridded-example.zarr/sea_ice_fraction/c/0/1/0 differ diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/sea_ice_fraction/c/0/1/1 b/beacon-arrow-zarr/test_files/gridded-example.zarr/sea_ice_fraction/c/0/1/1 new file mode 100644 index 00000000..5716f111 Binary files /dev/null and b/beacon-arrow-zarr/test_files/gridded-example.zarr/sea_ice_fraction/c/0/1/1 differ diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/sea_ice_fraction/c/0/2/0 b/beacon-arrow-zarr/test_files/gridded-example.zarr/sea_ice_fraction/c/0/2/0 new file mode 100644 index 00000000..5716f111 Binary files /dev/null and b/beacon-arrow-zarr/test_files/gridded-example.zarr/sea_ice_fraction/c/0/2/0 differ diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/sea_ice_fraction/c/0/2/1 b/beacon-arrow-zarr/test_files/gridded-example.zarr/sea_ice_fraction/c/0/2/1 new file mode 100644 index 00000000..5716f111 Binary files /dev/null and b/beacon-arrow-zarr/test_files/gridded-example.zarr/sea_ice_fraction/c/0/2/1 differ diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/sea_ice_fraction/c/0/3/0 b/beacon-arrow-zarr/test_files/gridded-example.zarr/sea_ice_fraction/c/0/3/0 new file mode 100644 index 00000000..5716f111 Binary files /dev/null and b/beacon-arrow-zarr/test_files/gridded-example.zarr/sea_ice_fraction/c/0/3/0 differ diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/sea_ice_fraction/c/0/3/1 b/beacon-arrow-zarr/test_files/gridded-example.zarr/sea_ice_fraction/c/0/3/1 new file mode 100644 index 00000000..5716f111 Binary files /dev/null and b/beacon-arrow-zarr/test_files/gridded-example.zarr/sea_ice_fraction/c/0/3/1 differ diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/sea_ice_fraction/zarr.json b/beacon-arrow-zarr/test_files/gridded-example.zarr/sea_ice_fraction/zarr.json new file mode 100644 index 00000000..46c7a7c3 --- /dev/null +++ b/beacon-arrow-zarr/test_files/gridded-example.zarr/sea_ice_fraction/zarr.json @@ -0,0 +1,59 @@ +{ + "shape": [ + 1, + 1208, + 1920 + ], + "data_type": "int8", + "chunk_grid": { + "name": "regular", + "configuration": { + "chunk_shape": [ + 1, + 302, + 960 + ] + } + }, + "chunk_key_encoding": { + "name": "default", + "configuration": { + "separator": "/" + } + }, + "fill_value": 0, + "codecs": [ + { + "name": "bytes" + }, + { + "name": "zstd", + "configuration": { + "level": 0, + "checksum": false + } + } + ], + "attributes": { + "standard_name": "sea_ice_area_fraction", + "long_name": "sea ice area fraction", + "units": "1", + "valid_min": 0, + "valid_max": 100, + "source_ice": "Not applicabile for Black sea", + "comment_ice": "Not applicabile for Black sea", + "source": "Not applicabile for Black sea", + "comment": "Not applicabile for Black sea", + "add_offset": 0.0, + "scale_factor": 0.009999999776482582, + "_FillValue": -128 + }, + "dimension_names": [ + "time", + "lat", + "lon" + ], + "zarr_format": 3, + "node_type": "array", + "storage_transformers": [] +} \ No newline at end of file diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/time/c/0 b/beacon-arrow-zarr/test_files/gridded-example.zarr/time/c/0 new file mode 100644 index 00000000..cda2c349 Binary files /dev/null and b/beacon-arrow-zarr/test_files/gridded-example.zarr/time/c/0 differ diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/time/zarr.json b/beacon-arrow-zarr/test_files/gridded-example.zarr/time/zarr.json new file mode 100644 index 00000000..848322b5 --- /dev/null +++ b/beacon-arrow-zarr/test_files/gridded-example.zarr/time/zarr.json @@ -0,0 +1,50 @@ +{ + "shape": [ + 1 + ], + "data_type": "int32", + "chunk_grid": { + "name": "regular", + "configuration": { + "chunk_shape": [ + 1 + ] + } + }, + "chunk_key_encoding": { + "name": "default", + "configuration": { + "separator": "/" + } + }, + "fill_value": 0, + "codecs": [ + { + "name": "bytes", + "configuration": { + "endian": "little" + } + }, + { + "name": "zstd", + "configuration": { + "level": 0, + "checksum": false + } + } + ], + "attributes": { + "long_name": "reference time of sst field", + "standard_name": "time", + "axis": "T", + "comment": "Nominal time of Level 4 analysis", + "units": "seconds since 1981-01-01", + "calendar": "Gregorian" + }, + "dimension_names": [ + "time" + ], + "zarr_format": 3, + "node_type": "array", + "storage_transformers": [] +} \ No newline at end of file diff --git a/beacon-arrow-zarr/test_files/gridded-example.zarr/zarr.json b/beacon-arrow-zarr/test_files/gridded-example.zarr/zarr.json new file mode 100644 index 00000000..15d08d59 --- /dev/null +++ b/beacon-arrow-zarr/test_files/gridded-example.zarr/zarr.json @@ -0,0 +1,454 @@ +{ + "attributes": { + "Conventions": "CF-1.4", + "title": "Black Sea SST Analysis, L4, 1km daily (SST_BS_SST_L4_NRT_OBSERVATIONS_010_006_c_V2)", + "summary": "Daily gap-free maps (L4) at 0.01deg. x 0.01deg. horizontal resolution over the Black Sea. The data (for day=dd) are obtained from infra-red measurements collected [since 7pm of day=dd-1 to 7am of day=dd] by satellite radiometers and statistical interpolation.It is the Copernicus sea surface temperature nominal operational product for the Black Sea in DT mode.", + "references": "Buongiorno Nardelli, B., C. Tronconi, a. Pisano, and R. Santoleri, 2013: High and Ultra-High resolution processing of satellite Sea Surface Temperature data over Southern European Seas in the framework of MyOcean project. /Remote Sens. Environ./, *129*, 1-16, doi:10.1016/j.rse.2012.10.012. AND Buongiorno Nardelli, B., a. Pisano, C. Tronconi, and R. Santoleri, 2015: Evaluation of different covariance models for the operational interpolation of high resolution satellite Sea Surface Temperature data over the Mediterranean Sea. /Remote Sens. Environ./, doi:10.1016/j.rse.2015.04.025.", + "institution": "GOS (Global Ocean Satellite monitoring and marine ecosystem study)", + "history": "GOS-MYO processor V1", + "comment": "WARNING: some applications are unable to properly handle byte values. If Values >127 are encounterd, please subtract 256", + "license": "free registration at Copernicus Marine Service (http://marine.copernicus.eu/web/56-user-registration-form.php)", + "id": "OISST_UHR_NRT-GOS-L4-BLK-v3.0", + "naming_authority": "org.ghrsst", + "product_version": "1.1", + "uuid": "f3079d2c-7d0f-11ee-8b9f-005056be6f27", + "gds_version_id": "v2.0.5", + "netcdf_version_id": "4.7.1, Build Date: OCT 6 2022", + "date_created": "20231107T014955Z", + "file_quality_level": 0, + "spatial_resolution": "0.01 degree", + "start_time": "20130228T190000Z", + "time_coverage_start": "20130228T190000Z", + "stop_time": "20130301T060000Z", + "time_coverage_end": "20130301T060000Z", + "software_version": "Copernicus Marine Service UHR L4 Processor V.1.0", + "westernmost_longitude": 26.379167556762695, + "easternmost_longitude": 42.37083435058594, + "southernmost_latitude": 38.75416564941406, + "northernmost_latitude": 48.8125, + "Scaling_Equation": "(scale_factor*data) + add_offset", + "source": "EUR-L2P-ATS_NR_2P, UPA-L2P-ATS_NR_2P, JPL-L2P-MODIS_A, JPL-L2P-MODIS_T, EUR-L2P-NAR17_SST, EUR-L2P-NAR18_SST, EUR-L2P-SEVIRI_SST, EUR-L2P-AVHRR_METOP_A, EUR-L2P-AVHRR_METOP_B, NAVO-L2P-VIIRS_NPP, EUR-L2P-SEVIRI_SST_Meteosat11, MAR-L2P-SLSTR_S3A, MAR-L2P-SLSTR_S3B, STAR-L2P-VIIRS_N20-ACSPO_V2.80", + "platform": "Envisat, NOAA-18, MetOpA, Aqua, Terra, MSG1, MetOpB, SUOMI-NPP, Meteosat_11, Sentinel 3A, Sentinel 3B, NOAA-20", + "sensor": "AATSR, SEVIRI, AVHRR, MODIS, VIIRS, SLSTR", + "Metadata_Conventions": "Unidata Dataset Discovery v1.0", + "metadata_link": "Link to collection metadata record at archive", + "keywords": "Oceans > Ocean Temperature > Sea Surface Temperature", + "keywords_vocabulary": "NASA Global Change Master Directory (GCMD) Science Keywords", + "standard_name_vocabulary": "NetCDF Climate and Forecast (CF) Metadata Convention", + "geospatial_lat_units": "degrees_north", + "geospatial_lat_resolution": 0.009999999776482582, + "geospatial_lon_units": "degrees_east", + "geospatial_lon_resolution": 0.009999999776482582, + "acknowledgment": "Please acknowledge the use of these data with the following statement: Generated/provided by Copernicus Marine Service and CNR - ISMAR ROME. We would also appreciate being informed of any publications.", + "creator_name": "ISMAR - Institute of Marine Sciences (CNR - Rome)", + "creator_email": "gsdk@artov.ismar.cnr.it", + "creator_url": "http://gosweb.artov.isac.cnr.it/", + "project": "Copernicus Marine Environment Monitoring Service (CMEMS)", + "publisher_name": "CNR ISMAR GOS - CMEMS SST-TAC", + "publisher_url": "http://marine.copernicus.eu/", + "publisher_email": "servicedesk.cmems@mercator-ocean.eu, gsdk@artov.ismar.cnr.it", + "processing_level": "L4", + "cdm_data_type": "grid" + }, + "zarr_format": 3, + "consolidated_metadata": { + "kind": "inline", + "must_understand": false, + "metadata": { + "time": { + "shape": [ + 1 + ], + "data_type": "int32", + "chunk_grid": { + "name": "regular", + "configuration": { + "chunk_shape": [ + 1 + ] + } + }, + "chunk_key_encoding": { + "name": "default", + "configuration": { + "separator": "/" + } + }, + "fill_value": 0, + "codecs": [ + { + "name": "bytes", + "configuration": { + "endian": "little" + } + }, + { + "name": "zstd", + "configuration": { + "level": 0, + "checksum": false + } + } + ], + "attributes": { + "long_name": "reference time of sst field", + "standard_name": "time", + "axis": "T", + "comment": "Nominal time of Level 4 analysis", + "units": "seconds since 1981-01-01", + "calendar": "Gregorian" + }, + "dimension_names": [ + "time" + ], + "zarr_format": 3, + "node_type": "array", + "storage_transformers": [] + }, + "mask": { + "shape": [ + 1, + 1208, + 1920 + ], + "data_type": "int8", + "chunk_grid": { + "name": "regular", + "configuration": { + "chunk_shape": [ + 1, + 302, + 960 + ] + } + }, + "chunk_key_encoding": { + "name": "default", + "configuration": { + "separator": "/" + } + }, + "fill_value": 0, + "codecs": [ + { + "name": "bytes" + }, + { + "name": "zstd", + "configuration": { + "level": 0, + "checksum": false + } + } + ], + "attributes": { + "long_name": "sea/land/lake/ice field composite mask", + "flag_values": [ + 1, + 2, + 3, + 4 + ], + "flag_meanings": "sea land lake ice", + "valid_min": 1, + "valid_max": 4, + "source_mask": "----", + "comment": "b0:1= sea; b1:1= land; b2:1= lake; b3:1= sea ice; b4-7: reserved for grid mask", + "_FillValue": -128 + }, + "dimension_names": [ + "time", + "lat", + "lon" + ], + "zarr_format": 3, + "node_type": "array", + "storage_transformers": [] + }, + "sea_ice_fraction": { + "shape": [ + 1, + 1208, + 1920 + ], + "data_type": "int8", + "chunk_grid": { + "name": "regular", + "configuration": { + "chunk_shape": [ + 1, + 302, + 960 + ] + } + }, + "chunk_key_encoding": { + "name": "default", + "configuration": { + "separator": "/" + } + }, + "fill_value": 0, + "codecs": [ + { + "name": "bytes" + }, + { + "name": "zstd", + "configuration": { + "level": 0, + "checksum": false + } + } + ], + "attributes": { + "standard_name": "sea_ice_area_fraction", + "long_name": "sea ice area fraction", + "units": "1", + "valid_min": 0, + "valid_max": 100, + "source_ice": "Not applicabile for Black sea", + "comment_ice": "Not applicabile for Black sea", + "source": "Not applicabile for Black sea", + "comment": "Not applicabile for Black sea", + "add_offset": 0.0, + "scale_factor": 0.009999999776482582, + "_FillValue": -128 + }, + "dimension_names": [ + "time", + "lat", + "lon" + ], + "zarr_format": 3, + "node_type": "array", + "storage_transformers": [] + }, + "lon": { + "shape": [ + 1920 + ], + "data_type": "float32", + "chunk_grid": { + "name": "regular", + "configuration": { + "chunk_shape": [ + 1920 + ] + } + }, + "chunk_key_encoding": { + "name": "default", + "configuration": { + "separator": "/" + } + }, + "fill_value": 0.0, + "codecs": [ + { + "name": "bytes", + "configuration": { + "endian": "little" + } + }, + { + "name": "zstd", + "configuration": { + "level": 0, + "checksum": false + } + } + ], + "attributes": { + "long_name": "longitude", + "standard_name": "longitude", + "axis": "X", + "units": "degrees_east", + "comment": "geographical coordinates, WGS84 projection", + "valid_min": 25.0, + "valid_max": 43.0, + "_FillValue": "AAAAAAAA+H8=" + }, + "dimension_names": [ + "lon" + ], + "zarr_format": 3, + "node_type": "array", + "storage_transformers": [] + }, + "analysed_sst": { + "shape": [ + 1, + 1208, + 1920 + ], + "data_type": "int16", + "chunk_grid": { + "name": "regular", + "configuration": { + "chunk_shape": [ + 1, + 302, + 960 + ] + } + }, + "chunk_key_encoding": { + "name": "default", + "configuration": { + "separator": "/" + } + }, + "fill_value": 0, + "codecs": [ + { + "name": "bytes", + "configuration": { + "endian": "little" + } + }, + { + "name": "zstd", + "configuration": { + "level": 0, + "checksum": false + } + } + ], + "attributes": { + "long_name": "analysed sea surface temperature", + "standard_name": "sea_surface_foundation_temperature", + "type": "foundation", + "units": "kelvin", + "valid_min": -300, + "valid_max": 4500, + "source": "EUR-L2P-ATS_NR_2P,UPA-L2P-ATS_NR_2P,JPL-L2P-MODIS_A,JPL-L2P-MODIS_T,EUR-L2P-NAR17_SST,EUR-L2P-NAR18_SST,EUR-L2P-SEVIRI_SST,EUR-L2P-AVHRR_METOP_A,EUR-L2P-AVHRR_METOP_B,NAVO-L2P-VIIRS_NPP,EUR-L2P-SEVIRI_SST_Meteosat11,MAR-L2P-SLSTR_S3A,MAR-L2P-SLSTR_S3B,STAR-L2P-VIIRS_N20-ACSPO_V2.80", + "comment": "Optimal interpolation (OI) at Ultrahigh resolution of sst measurements from infrared sensors", + "add_offset": 273.1499938964844, + "scale_factor": 0.009999999776482582, + "_FillValue": -32768 + }, + "dimension_names": [ + "time", + "lat", + "lon" + ], + "zarr_format": 3, + "node_type": "array", + "storage_transformers": [] + }, + "lat": { + "shape": [ + 1208 + ], + "data_type": "float32", + "chunk_grid": { + "name": "regular", + "configuration": { + "chunk_shape": [ + 1208 + ] + } + }, + "chunk_key_encoding": { + "name": "default", + "configuration": { + "separator": "/" + } + }, + "fill_value": 0.0, + "codecs": [ + { + "name": "bytes", + "configuration": { + "endian": "little" + } + }, + { + "name": "zstd", + "configuration": { + "level": 0, + "checksum": false + } + } + ], + "attributes": { + "long_name": "latitude", + "standard_name": "latitude", + "axis": "Y", + "units": "degrees_north", + "comment": "geographical coordinates, WGS84 projection", + "valid_min": 37.0, + "valid_max": 49.0, + "_FillValue": "AAAAAAAA+H8=" + }, + "dimension_names": [ + "lat" + ], + "zarr_format": 3, + "node_type": "array", + "storage_transformers": [] + }, + "analysis_error": { + "shape": [ + 1, + 1208, + 1920 + ], + "data_type": "int16", + "chunk_grid": { + "name": "regular", + "configuration": { + "chunk_shape": [ + 1, + 302, + 960 + ] + } + }, + "chunk_key_encoding": { + "name": "default", + "configuration": { + "separator": "/" + } + }, + "fill_value": 0, + "codecs": [ + { + "name": "bytes", + "configuration": { + "endian": "little" + } + }, + { + "name": "zstd", + "configuration": { + "level": 0, + "checksum": false + } + } + ], + "attributes": { + "long_name": "estimated error standard deviation of analysed_sst", + "units": "kelvin", + "valid_min": 0, + "valid_max": 10000, + "source": "EUR-L2P-ATS_NR_2P,UPA-L2P-ATS_NR_2P,JPL-L2P-MODIS_A,JPL-L2P-MODIS_T,EUR-L2P-NAR17_SST,EUR-L2P-NAR18_SST,EUR-L2P-SEVIRI_SST,EUR-L2P-AVHRR_METOP_A,EUR-L2P-AVHRR_METOP_B,NAVO-L2P-VIIRS_NPP,EUR-L2P-SEVIRI_SST_Meteosat11,MAR-L2P-SLSTR_S3A,MAR-L2P-SLSTR_S3B,STAR-L2P-VIIRS_N20-ACSPO_V2.80", + "comment": "Error obtained by Optimal interpolation (OI) at Ultrahigh resolution of sst measurements from infrared sensors", + "add_offset": 0.0, + "scale_factor": 0.009999999776482582, + "_FillValue": -32768 + }, + "dimension_names": [ + "time", + "lat", + "lon" + ], + "zarr_format": 3, + "node_type": "array", + "storage_transformers": [] + } + } + }, + "node_type": "group" +} \ No newline at end of file diff --git a/beacon-data-lake/Cargo.toml b/beacon-data-lake/Cargo.toml index 6ab912c2..6aeea5de 100644 --- a/beacon-data-lake/Cargo.toml +++ b/beacon-data-lake/Cargo.toml @@ -35,4 +35,5 @@ utoipa = { workspace=true} beacon-config = { path = "../beacon-config" } -beacon-formats = { path = "../beacon-formats" } \ No newline at end of file +beacon-formats = { path = "../beacon-formats" } +beacon-arrow-zarr = { path = "../beacon-arrow-zarr" } \ No newline at end of file diff --git a/beacon-data-lake/src/lib.rs b/beacon-data-lake/src/lib.rs index f175adf1..98ddedf4 100644 --- a/beacon-data-lake/src/lib.rs +++ b/beacon-data-lake/src/lib.rs @@ -17,12 +17,11 @@ use datafusion::{ }; use futures::StreamExt; use object_store::{ObjectStore, aws::AmazonS3Builder, local::LocalFileSystem, path::PathPart}; -use url::Url; use crate::{ files::{collection::FileCollection, temp_output_file::TempOutputFile}, table::{Table, empty::EmptyTable, error::TableError}, - util::split_glob, + util::parse_listing_table_url, }; pub mod files; @@ -72,50 +71,16 @@ impl Debug for DataLake { } impl DataLake { + #[inline(always)] pub fn try_create_listing_url( &self, path: String, ) -> datafusion::error::Result { - let parts = self.data_directory_prefix.parts().collect::>(); - if let Some((base, pattern)) = split_glob(&path) { - let pattern = glob::Pattern::new(&pattern).unwrap(); - let mut full_path = format!( - "{}{}", - self.data_directory_store_url, - object_store::path::Path::from_iter(parts.clone().into_iter()), - ); - if base.components().next().is_some() { - full_path.push_str(format!("/{}", base.as_os_str().to_string_lossy()).as_str()); - } - if !full_path.ends_with('/') { - full_path.push('/'); - } - - let url = Url::parse(&full_path).unwrap(); - - let table_url = ListingTableUrl::try_new(url, Some(pattern))?; - - Ok(table_url) - } else { - let mut full_path = format!( - "{}{}", - self.data_directory_store_url, - object_store::path::Path::from_iter(parts.clone().into_iter()), - ); - - // If full path ends with '/' then just append the path, otherwise append without a '/' - if full_path.ends_with('/') { - full_path.push_str(&path); - } else { - full_path.push_str(format!("/{}", path).as_str()); - } - - let url = Url::parse(&full_path).unwrap(); - - let table_url = ListingTableUrl::try_new(url, None)?; - - Ok(table_url) - } + parse_listing_table_url( + &self.data_directory_store_url, + &self.data_directory_prefix, + &path, + ) } pub fn netcdf_object_resolver() -> Arc { diff --git a/beacon-data-lake/src/table/error.rs b/beacon-data-lake/src/table/error.rs index 9a4d3185..109989cd 100644 --- a/beacon-data-lake/src/table/error.rs +++ b/beacon-data-lake/src/table/error.rs @@ -10,4 +10,6 @@ pub enum TableError { TableAlreadyExists(String), #[error("Table not found: {0}")] TableNotFound(String), + #[error("Datafusion Error: {0}")] + DatafusionError(#[from] datafusion::error::DataFusionError), } diff --git a/beacon-data-lake/src/table/logical/mod.rs b/beacon-data-lake/src/table/logical/mod.rs index 4208858c..b04be813 100644 --- a/beacon-data-lake/src/table/logical/mod.rs +++ b/beacon-data-lake/src/table/logical/mod.rs @@ -1,16 +1,13 @@ use std::sync::Arc; use datafusion::{ - catalog::TableProvider, datasource::listing::ListingTableUrl, - execution::object_store::ObjectStoreUrl, prelude::SessionContext, + catalog::TableProvider, execution::object_store::ObjectStoreUrl, prelude::SessionContext, }; -use object_store::path::PathPart; -use url::Url; use crate::{ files::collection::FileCollection, table::{error::TableError, table_formats::TableFileFormat}, - util::split_glob, + util::parse_listing_table_url, }; /// A representation of a logical table configuration. @@ -44,63 +41,14 @@ impl LogicalTable { let session_state = session_context.state(); let mut table_urls = Vec::new(); - let parts: Vec> = data_directory_prefix - .parts() - .map(|part| part.as_ref().to_string().into()) - .collect::>(); // Construct table URLs based on the paths provided in the logical table. - for path in &self.glob_paths { - // Check if path contains glob expression - if let Some((base, pattern)) = split_glob(path) { - let pattern = glob::Pattern::new(&pattern).map_err(|e| { - TableError::GenericTableError(format!("Failed to parse glob pattern: {}", e)) - })?; - let mut full_path = format!( - "{}{}", - data_directory_store_url, - object_store::path::Path::from_iter(parts.clone().into_iter()), - ); - if base.components().next().is_some() { - full_path.push_str(format!("/{}", base.as_os_str().to_string_lossy()).as_str()); - } - if !full_path.ends_with('/') { - full_path.push('/'); - } - - let url = Url::parse(&full_path).map_err(|e| { - TableError::GenericTableError(format!("Failed to parse URL: {}", e)) - })?; - - let table_url = ListingTableUrl::try_new(url, Some(pattern)).map_err(|e| { - TableError::GenericTableError(format!("Failed to create table URL: {}", e)) - })?; - - table_urls.push(table_url); - } else { - let mut full_path = format!( - "{}{}", - data_directory_store_url, - object_store::path::Path::from_iter(parts.clone().into_iter()), - ); - - // If full path ends with '/' then just append the path, otherwise append without a '/' - if full_path.ends_with('/') { - full_path.push_str(path); - } else { - full_path.push_str(format!("/{}", path).as_str()); - } - - let url = Url::parse(&full_path).map_err(|e| { - TableError::GenericTableError(format!("Failed to parse URL: {}", e)) - })?; - - let table_url = ListingTableUrl::try_new(url, None).map_err(|e| { - TableError::GenericTableError(format!("Failed to create table URL: {}", e)) - })?; - - table_urls.push(table_url); - } + for glob_path in &self.glob_paths { + table_urls.push(parse_listing_table_url( + data_directory_store_url, + data_directory_prefix, + glob_path, + )?); } tracing::debug!( @@ -117,11 +65,17 @@ impl LogicalTable { )?; // Create the data source with the given file format and table URLs. - let source = FileCollection::new(&session_state, file_format_factory.default(), table_urls) - .await - .map_err(|e| { - TableError::GenericTableError(format!("Failed to create file collection: {}", e)) - })?; + let source = FileCollection::new( + &session_state, + self.file_format + .file_format() + .unwrap_or(file_format_factory.default()), + table_urls, + ) + .await + .map_err(|e| { + TableError::GenericTableError(format!("Failed to create file collection: {}", e)) + })?; Ok(Arc::new(source)) } diff --git a/beacon-data-lake/src/table/table_formats.rs b/beacon-data-lake/src/table/table_formats.rs index a7a5a0b6..549dcb32 100644 --- a/beacon-data-lake/src/table/table_formats.rs +++ b/beacon-data-lake/src/table/table_formats.rs @@ -1,10 +1,16 @@ -use std::sync::Arc; +use std::{collections::HashMap, sync::Arc}; -use beacon_formats::{arrow::ArrowFormat, csv::CsvFormat, parquet::ParquetFormat}; +use beacon_formats::zarr::{ + ZarrFormat, array_step_span::NumericArrayStepSpan, pushdown_statistics::ZarrPushDownStatistics, +}; +use datafusion::datasource::file_format::FileFormat; #[typetag::serde(tag = "file_format")] pub trait TableFileFormat: std::fmt::Debug + Send + Sync { fn file_ext(&self) -> String; + fn file_format(&self) -> Option> { + None + } } #[derive(Debug, serde::Deserialize, serde::Serialize)] @@ -48,3 +54,26 @@ impl TableFileFormat for NetCDFFileFormat { "nc".to_string() } } + +#[derive(Debug, serde::Deserialize, serde::Serialize)] +pub struct ZarrFileFormat { + #[serde(default)] + pub global_array_steps: HashMap, + #[serde(default)] + pub pushdown_statistics: ZarrPushDownStatistics, +} + +#[typetag::serde(name = "zarr")] +impl TableFileFormat for ZarrFileFormat { + fn file_ext(&self) -> String { + "zarr".to_string() + } + + fn file_format(&self) -> Option> { + Some(Arc::new( + ZarrFormat::default() + .with_array_steps(self.global_array_steps.clone()) + .with_pushdown_statistics(self.pushdown_statistics.clone()), + )) + } +} diff --git a/beacon-data-lake/src/util.rs b/beacon-data-lake/src/util.rs index 21a0b715..28181b69 100644 --- a/beacon-data-lake/src/util.rs +++ b/beacon-data-lake/src/util.rs @@ -1,16 +1,16 @@ -use std::{ - collections::HashMap, - path::{Path, PathBuf}, -}; +use std::{collections::HashMap, path::PathBuf}; use arrow::datatypes::{DataType, Fields, Schema, TimeUnit}; use datafusion::{ common::{ - Column, + Column, exec_datafusion_err, tree_node::{Transformed, TreeNode}, }, + datasource::listing::ListingTableUrl, + execution::object_store::ObjectStoreUrl, prelude::Expr, }; +use url::Url; #[derive(Debug, Clone, thiserror::Error)] pub enum SuperTypeError { @@ -293,17 +293,160 @@ pub fn remap_filter( Ok(transformed.data) } -pub fn split_glob(path: &str) -> Option<(PathBuf, String)> { - let p = Path::new(path); +pub fn parse_listing_table_url( + data_directory_store_url: &ObjectStoreUrl, + data_directory_prefix: &object_store::path::Path, + glob_path: &str, +) -> datafusion::error::Result { + // Check if path contains glob expression + let (path, glob_pattern) = split_path_and_glob(glob_path); + let mut full_path = format!("{}{}", data_directory_store_url, data_directory_prefix); + if path.components().next().is_some() { + full_path.push_str(format!("/{}", path.as_os_str().to_string_lossy()).as_str()); + } + if !full_path.ends_with('/') { + full_path.push('/'); + } - // Get the last component (file name or pattern) - if let Some(last) = p.file_name().and_then(|os| os.to_str()) { - // Does it look like a glob? - if last.contains('*') || last.contains('?') || last.contains('[') { - // Base path = everything before the last component - let base = p.parent().unwrap_or(Path::new("")).to_path_buf(); - return Some((base, last.to_string())); - } + let glob_pattern_parsed = if let Some(pattern) = glob_pattern { + Some( + glob::Pattern::new(&pattern) + .map_err(|e| exec_datafusion_err!("Failed to parse glob pattern: {}", e))?, + ) + } else { + None + }; + + let url = + Url::parse(&full_path).map_err(|e| exec_datafusion_err!("Failed to parse URL: {}", e))?; + + let table_url = ListingTableUrl::try_new(url, glob_pattern_parsed) + .map_err(|e| exec_datafusion_err!("Failed to create table URL: {}", e))?; + + Ok(table_url) +} + +/// Splits a path containing an optional glob pattern into +/// (base_path, optional_glob_pattern). +/// +/// Examples: +/// - "src/**/*.rs" -> ("src", Some("**/*.rs")) +/// - "src/*.rs" -> ("src", Some("*.rs")) +/// - "*.rs" -> (".", Some("*.rs")) +/// - "src/main.rs" -> ("src/main.rs", None) +pub fn split_path_and_glob(input: &str) -> (PathBuf, Option) { + let glob_chars = ['*', '?', '[', ']']; + + // Split the path into components manually (cross-platform) + let parts: Vec<&str> = input.split(['/', '\\']).collect(); + + // Find the index of the first segment that contains any glob char + if let Some(i) = parts + .iter() + .position(|p| p.chars().any(|c| glob_chars.contains(&c))) + { + let base = if i == 0 { + PathBuf::from(".") + } else { + PathBuf::from(parts[..i].join(std::path::MAIN_SEPARATOR_STR)) + }; + let glob = parts[i..].join(std::path::MAIN_SEPARATOR_STR); + (base, Some(glob)) + } else { + // No glob pattern at all + (PathBuf::from(input), None) + } +} + +#[cfg(test)] +mod tests { + use std::path::MAIN_SEPARATOR; + + use super::*; + + #[test] + fn test_no_glob() { + let (base, glob) = split_path_and_glob("src/main.rs"); + assert_eq!(base, PathBuf::from("src/main.rs")); + assert_eq!(glob, None); + } + + #[test] + fn test_simple_glob() { + let (base, glob) = split_path_and_glob("src/*.rs"); + assert_eq!(base, PathBuf::from("src")); + assert_eq!(glob, Some("*.rs".to_string())); + } + + #[test] + fn test_recursive_glob() { + let (base, glob) = split_path_and_glob("src/**/*.rs"); + assert_eq!(base, PathBuf::from("src")); + assert_eq!(glob, Some(format!("**{}*.rs", std::path::MAIN_SEPARATOR))); + } + + #[test] + fn test_glob_at_start() { + let (base, glob) = split_path_and_glob("*.txt"); + assert_eq!(base, PathBuf::from(".")); + assert_eq!(glob, Some("*.txt".to_string())); + } + + #[test] + fn test_glob_in_middle_segment() { + let (base, glob) = split_path_and_glob("foo[ab]/bar?.txt"); + assert_eq!(base, PathBuf::from("foo[ab]")); + assert_eq!(glob, Some("bar?.txt".to_string())); + } + + #[test] + fn test_no_separator_but_glob_present() { + let (base, glob) = split_path_and_glob("file[0-9].log"); + assert_eq!(base, PathBuf::from(".")); + assert_eq!(glob, Some("file[0-9].log".to_string())); + } + + #[test] + fn test_windows_style_path() { + // Works on both Unix and Windows separators + let input = format!("dir{}**{}*.txt", MAIN_SEPARATOR, MAIN_SEPARATOR); + let (base, glob) = split_path_and_glob(&input); + assert_eq!(base, PathBuf::from("dir")); + assert_eq!(glob, Some(format!("**{}*.txt", MAIN_SEPARATOR))); + } + + #[test] + fn test_parse_listing_table_url_no_glob() { + use datafusion::execution::object_store::ObjectStoreUrl; + + let store_url: ObjectStoreUrl = ObjectStoreUrl::parse("file:///").unwrap(); + let prefix = object_store::path::Path::from("datasets"); + let url = parse_listing_table_url(&store_url, &prefix, "foo/zarr.json").unwrap(); + // Debug representation should contain the base path we constructed + let dbg = format!("{:?}", url); + assert!(dbg.contains("/datasets/foo/"), "dbg={}", dbg); + } + + #[test] + fn test_parse_listing_table_url_with_glob() { + use datafusion::execution::object_store::ObjectStoreUrl; + + let store_url: ObjectStoreUrl = ObjectStoreUrl::parse("file:///").unwrap(); + let prefix = object_store::path::Path::from("datasets"); + let url = parse_listing_table_url(&store_url, &prefix, "subdir/*.json").unwrap(); + let dbg = format!("{:?}", url); + assert!(dbg.contains("/datasets/subdir/"), "dbg={}", dbg); + // pattern should be present in debug output + assert!(dbg.contains("*.json"), "dbg={}", dbg); + } + + #[test] + fn test_parse_listing_table_url_invalid_glob() { + use datafusion::execution::object_store::ObjectStoreUrl; + + let store_url: ObjectStoreUrl = ObjectStoreUrl::parse("file:///").unwrap(); + let prefix = object_store::path::Path::from("datasets"); + let res = parse_listing_table_url(&store_url, &prefix, "subdir/[invalid"); + assert!(res.is_err(), "expected error for invalid glob pattern"); } - None } diff --git a/beacon-formats/Cargo.toml b/beacon-formats/Cargo.toml index 143e611a..54ddbfde 100644 --- a/beacon-formats/Cargo.toml +++ b/beacon-formats/Cargo.toml @@ -25,7 +25,16 @@ tokio = { workspace = true } async_zip = { version = "0.0.18", features = ["tokio", "zstd", "deflate"] } tokio-util = "0.7.16" +zarrs_object_store = "0.5.0" +zarrs_storage = "0.4.0" +zarrs = {"version" = "0.22.2", features = ["async"]} + +url = "2.5.4" +glob = "0.3.3" +chrono = { workspace = true } beacon-common = { path = "../beacon-common" } beacon-arrow-netcdf = { path = "../beacon-arrow-netcdf" } -beacon-arrow-odv = { path = "../beacon-arrow-odv" } \ No newline at end of file +beacon-arrow-odv = { path = "../beacon-arrow-odv" } +beacon-arrow-zarr = { path = "../beacon-arrow-zarr" } +nd-arrow-array = { git = "https://github.com/maris-development/nd-arrow-array.git", branch = "main" } \ No newline at end of file diff --git a/beacon-formats/src/lib.rs b/beacon-formats/src/lib.rs index f8f8fea2..51a916fe 100644 --- a/beacon-formats/src/lib.rs +++ b/beacon-formats/src/lib.rs @@ -11,6 +11,7 @@ use crate::{ object_resolver::{NetCDFObjectResolver, NetCDFSinkResolver}, }, parquet::ParquetFormatFactory, + zarr::ZarrFormatFactory, }; pub mod arrow; @@ -19,6 +20,7 @@ pub mod geo_parquet; pub mod netcdf; pub mod odv_ascii; pub mod parquet; +pub mod zarr; /// Register file formats with the session state that can be used for reading pub fn register_file_formats( @@ -38,5 +40,6 @@ pub fn register_file_formats( )), true, )?; + session_state.register_file_format(Arc::new(ZarrFormatFactory), true)?; Ok(()) } diff --git a/beacon-formats/src/zarr/array_step_span.rs b/beacon-formats/src/zarr/array_step_span.rs new file mode 100644 index 00000000..5a993c44 --- /dev/null +++ b/beacon-formats/src/zarr/array_step_span.rs @@ -0,0 +1,8 @@ +#[derive(Clone, Debug, serde::Deserialize, serde::Serialize)] +pub struct NumericArrayStepSpan { + pub dimension: String, + pub column: String, + pub start: f64, + pub end: f64, + pub step: f64, +} diff --git a/beacon-formats/src/zarr/expr_util.rs b/beacon-formats/src/zarr/expr_util.rs new file mode 100644 index 00000000..ce665b2a --- /dev/null +++ b/beacon-formats/src/zarr/expr_util.rs @@ -0,0 +1,232 @@ +use datafusion::logical_expr::Operator; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_expr::expressions::{BinaryExpr, CastExpr, Column, Literal}; +use datafusion::scalar::ScalarValue; +use std::sync::Arc; + +/// A single interval: (value, inclusive?) +#[derive(Debug, Clone, Default)] +pub struct Range { + pub min: Option<(ScalarValue, bool)>, + pub max: Option<(ScalarValue, bool)>, +} + +impl Range { + pub fn as_f64_max(&self) -> Option { + self.max.as_ref().and_then(|(v, _)| { + if let ScalarValue::Float64(Some(f)) = v { + Some(*f) + } else { + None + } + }) + } + + pub fn as_f64_min(&self) -> Option { + self.min.as_ref().and_then(|(v, _)| { + if let ScalarValue::Float64(Some(f)) = v { + Some(*f) + } else { + None + } + }) + } + + pub fn as_timestamp_max(&self) -> Option { + self.max.as_ref().and_then(|(v, _)| match v { + ScalarValue::TimestampNanosecond(Some(ts), None) => { + chrono::NaiveDateTime::from_timestamp_nanos(*ts) + } + ScalarValue::TimestampMicrosecond(Some(ts), None) => { + chrono::NaiveDateTime::from_timestamp_micros(*ts) + } + ScalarValue::TimestampMillisecond(Some(ts), None) => { + chrono::NaiveDateTime::from_timestamp_millis(*ts) + } + ScalarValue::TimestampSecond(Some(ts), None) => { + Some(chrono::NaiveDateTime::from_timestamp(*ts, 0)) + } + _ => None, + }) + } + + pub fn as_timestamp_min(&self) -> Option { + self.min.as_ref().and_then(|(v, _)| match v { + ScalarValue::TimestampNanosecond(Some(ts), None) => { + chrono::NaiveDateTime::from_timestamp_nanos(*ts) + } + ScalarValue::TimestampMicrosecond(Some(ts), None) => { + chrono::NaiveDateTime::from_timestamp_micros(*ts) + } + ScalarValue::TimestampMillisecond(Some(ts), None) => { + chrono::NaiveDateTime::from_timestamp_millis(*ts) + } + ScalarValue::TimestampSecond(Some(ts), None) => { + Some(chrono::NaiveDateTime::from_timestamp(*ts, 0)) + } + _ => None, + }) + } + + pub fn is_max_inclusive(&self) -> Option { + self.max.as_ref().map(|(_, inc)| *inc) + } + pub fn is_min_inclusive(&self) -> Option { + self.min.as_ref().map(|(_, inc)| *inc) + } + + fn tighten_lower(&mut self, val: ScalarValue, inclusive: bool) { + match &self.min { + None => self.min = Some((val, inclusive)), + Some((cur, cur_inc)) => { + if let Some(ord) = cur.partial_cmp(&val) { + use std::cmp::Ordering::*; + match ord { + Less => self.min = Some((val, inclusive)), + Equal => self.min = Some((val, *cur_inc && inclusive)), + Greater => {} // keep tighter existing lower + } + } + } + } + } + fn tighten_upper(&mut self, val: ScalarValue, inclusive: bool) { + match &self.max { + None => self.max = Some((val, inclusive)), + Some((cur, cur_inc)) => { + if let Some(ord) = cur.partial_cmp(&val) { + use std::cmp::Ordering::*; + match ord { + Greater => self.max = Some((val, inclusive)), + Equal => self.max = Some((val, *cur_inc && inclusive)), + Less => {} // keep tighter existing upper + } + } + } + } + } +} + +/// Entry point: extract tightest range for `target_col` from a set of +/// physical filter expressions (possibly complex AND trees). +pub fn extract_range_from_physical_filters( + filters: &[Arc], + target_col: &str, +) -> Option { + let mut out = Range::default(); + for f in filters { + walk_physical_expr(f, target_col, &mut out); + } + + if out.min.is_some() || out.max.is_some() { + Some(out) + } else { + None + } +} + +fn walk_physical_expr(node: &Arc, target_col: &str, out: &mut Range) { + if let Some(bin) = downcast::(node) { + let op = bin.op(); + + // AND = BinaryExpr with Operator::And: visit both sides + if *op == Operator::And { + walk_physical_expr(bin.left(), target_col, out); + walk_physical_expr(bin.right(), target_col, out); + return; + } + + // Comparison: try (col vs lit), then (lit vs col) with flipped op + if let (Some(_c), Some(lv)) = ( + as_target_col(bin.left(), target_col), + as_literal(bin.right()), + ) { + apply_cmp(out, *op, lv); + return; + } + if let (Some(lv), Some(_c)) = ( + as_literal(bin.left()), + as_target_col(bin.right(), target_col), + ) { + apply_cmp(out, flip_op(*op), lv); + return; + } + + // Not a direct pattern: still walk children (there might be nested pieces) + walk_physical_expr(bin.left(), target_col, out); + walk_physical_expr(bin.right(), target_col, out); + return; + } + + // Generic: peel through any children (e.g., some wrappers/functions you may want to inspect) + for child in node.children() { + walk_physical_expr(child, target_col, out); + } +} + +fn apply_cmp(out: &mut Range, op: Operator, lit: &ScalarValue) { + match op { + Operator::Gt => out.tighten_lower(lit.clone(), false), + Operator::GtEq => out.tighten_lower(lit.clone(), true), + Operator::Lt => out.tighten_upper(lit.clone(), false), + Operator::LtEq => out.tighten_upper(lit.clone(), true), + Operator::Eq => { + out.tighten_lower(lit.clone(), true); + out.tighten_upper(lit.clone(), true); + } + _ => {} // ignore OR, NotEq, etc. (not a single-interval constraint) + } +} + +fn flip_op(op: Operator) -> Operator { + match op { + Operator::Lt => Operator::Gt, + Operator::LtEq => Operator::GtEq, + Operator::Gt => Operator::Lt, + Operator::GtEq => Operator::LtEq, + other => other, + } +} + +/// ---- helpers ---- + +fn as_literal(e: &Arc) -> Option<&ScalarValue> { + match peel_casts(e) { + Some(l) if downcast::(l).is_some() => { + downcast::(l).map(|lit| lit.value()) + } + _ => None, + } +} + +fn as_target_col<'l>(e: &'l Arc, target: &str) -> Option<&'l Column> { + match peel_casts(e) { + Some(c) if downcast::(c).is_some() => { + let c = downcast::(c).unwrap(); + // Column::name() is the leaf name at physical layer; qualify if you pass "table.col" + if c.name() == target || format!("{}", c).ends_with(&format!(".{}", target)) { + Some(c) + } else { + None + } + } + _ => None, + } +} + +/// Remove any number of CAST(...) wrappers around an expression. +fn peel_casts(mut e: &Arc) -> Option<&Arc> { + let mut cur = e; + loop { + if let Some(c) = downcast::(cur) { + cur = c.expr(); + continue; + } + return Some(cur); + } +} + +/// Convenience downcast +fn downcast(expr: &Arc) -> Option<&T> { + expr.as_any().downcast_ref::() +} diff --git a/beacon-formats/src/zarr/mod.rs b/beacon-formats/src/zarr/mod.rs new file mode 100644 index 00000000..48fea513 --- /dev/null +++ b/beacon-formats/src/zarr/mod.rs @@ -0,0 +1,358 @@ +use std::{ + collections::{HashMap, HashSet}, + fmt::Debug, + sync::Arc, +}; + +use arrow::datatypes::SchemaRef; +use beacon_arrow_zarr::reader::AsyncArrowZarrGroupReader; +use datafusion::{ + catalog::{Session, memory::DataSourceExec}, + common::{ColumnStatistics, GetExt, Statistics}, + datasource::{ + file_format::{FileFormat, FileFormatFactory, file_compression_type::FileCompressionType}, + listing::PartitionedFile, + physical_plan::{ + FileGroup, FileGroupPartitioner, FileScanConfig, FileScanConfigBuilder, FileSource, + }, + }, + physical_plan::ExecutionPlan, +}; +use object_store::{ObjectMeta, ObjectStore}; +use zarrs::group::Group; +use zarrs_object_store::AsyncObjectStore; +use zarrs_storage::AsyncReadableListableStorageTraits; + +use crate::zarr::{ + array_step_span::NumericArrayStepSpan, + pushdown_statistics::ZarrPushDownStatistics, + source::{ZarrSource, fetch_schema}, + util::{ZarrPath, is_zarr_v3_metadata, path_parent, top_level_zarr_meta_v3}, +}; + +pub mod array_step_span; +pub mod expr_util; +pub mod opener; +mod partition; +pub mod pushdown_statistics; +mod source; +mod stream_share; +pub mod util; + +pub struct ZarrFormatFactory; + +impl Debug for ZarrFormatFactory { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ZarrFormatFactory").finish() + } +} + +impl ZarrFormatFactory { + pub fn new() -> Self { + Self + } +} + +impl Default for ZarrFormatFactory { + fn default() -> Self { + Self::new() + } +} + +impl GetExt for ZarrFormatFactory { + fn get_ext(&self) -> String { + "zarr".to_string() + } +} + +impl FileFormatFactory for ZarrFormatFactory { + fn create( + &self, + state: &dyn datafusion::catalog::Session, + format_options: &std::collections::HashMap, + ) -> datafusion::error::Result> { + // Here you can use `state` and `format_options` to customize the creation of the FileFormat + Ok(Arc::new(ZarrFormat::default())) + } + + fn default(&self) -> std::sync::Arc { + Arc::new(ZarrFormat::default()) + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } +} + +#[derive(Debug, Clone, Default)] +pub struct ZarrFormat { + array_steps: HashMap, + zarr_pushdown_statistics: ZarrPushDownStatistics, +} + +impl ZarrFormat { + pub fn with_array_steps(mut self, array_steps: HashMap) -> Self { + self.array_steps = array_steps; + self + } + + pub fn with_pushdown_statistics( + mut self, + zarr_pushdown_statistics: ZarrPushDownStatistics, + ) -> Self { + self.zarr_pushdown_statistics = zarr_pushdown_statistics; + self + } +} + +#[async_trait::async_trait] +impl FileFormat for ZarrFormat { + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn compression_type(&self) -> Option { + None + } + + fn get_ext(&self) -> String { + "zarr.json".to_string() + } + + fn get_ext_with_compression( + &self, + _file_compression_type: &FileCompressionType, + ) -> datafusion::error::Result { + Ok("zarr.json".to_string()) + } + + async fn infer_schema( + &self, + _state: &dyn Session, + store: &Arc, + objects: &[ObjectMeta], + ) -> datafusion::error::Result { + // Ensure every group is a zarr.json file. + for object in objects { + if !is_zarr_v3_metadata(object) { + return Err(datafusion::error::DataFusionError::Execution(format!( + "Object at location '{}' is not a Zarr v3 metadata file (zarr.json)", + object.location.as_ref() + ))); + } + } + // Collect only top-level zarr.json files. + let verified_objects = top_level_zarr_meta_v3(objects); + tracing::debug!("Top-level Zarr v3 groups: {:?}", verified_objects); + let mut schema = None; + for object in verified_objects { + let zarr_path = ZarrPath::new_from_object_meta(object.clone()).map_err(|e| { + datafusion::error::DataFusionError::Execution(format!( + "Failed to create ZarrPath from ObjectMeta at {}: {}", + object.location, e + )) + })?; + let object_schema = fetch_schema(store.clone(), &zarr_path).await?; + schema = Some(object_schema); + } + Ok(schema.unwrap()) + } + + async fn infer_stats( + &self, + _state: &dyn Session, + _store: &Arc, + table_schema: SchemaRef, + _object: &ObjectMeta, + ) -> datafusion::error::Result { + Ok(Statistics::new_unknown(&table_schema)) + } + + async fn create_physical_plan( + &self, + state: &dyn Session, + conf: FileScanConfig, + ) -> datafusion::error::Result> { + // Recreate file groups here. Only maintain top-level groups. + let mut object_metas: Vec = Vec::new(); + for group in &conf.file_groups { + for file in group.files() { + object_metas.push(file.object_meta.clone()); + } + } + let object_store = state + .runtime_env() + .object_store(conf.object_store_url.clone())?; + + // Recreate file groups with only top-level Zarr groups. + let top_level_metas = top_level_zarr_meta_v3(&object_metas); + let mut file_groups: Vec = vec![]; + for meta in top_level_metas { + let file = self + .partition_zarr_group(&conf.file_schema, &meta, object_store.clone()) + .await?; + file_groups.push(file); + } + + let source = ZarrSource::default() + .with_array_steps(self.array_steps.clone()) + .with_pushdown_zarr_statistics(self.zarr_pushdown_statistics.clone()); + let conf = FileScanConfigBuilder::from(conf) + .with_file_groups(file_groups) + .with_source(Arc::new(source)) + .build(); + Ok(DataSourceExec::from_data_source(conf)) + } + + fn file_source(&self) -> Arc { + Arc::new( + ZarrSource::default() + .with_array_steps(self.array_steps.clone()) + .with_pushdown_zarr_statistics(self.zarr_pushdown_statistics.clone()), + ) + } +} + +impl ZarrFormat { + async fn partition_zarr_group( + &self, + table_schema: &SchemaRef, + object: &ObjectMeta, + object_store: Arc, + ) -> datafusion::error::Result { + // The object can be assumed to be a top-level zarr.json file. + let zarr_store = Arc::new(AsyncObjectStore::new(object_store)) + as Arc; + let group_path = ZarrPath::new_from_object_meta(object.clone()).map_err(|e| { + datafusion::error::DataFusionError::Execution(format!( + "Failed to create ZarrPath from ObjectMeta at {}: {}", + object.location, e + )) + })?; + let group = Group::async_open(zarr_store, &group_path.as_zarr_path()) + .await + .map_err(|e| { + datafusion::error::DataFusionError::Execution(format!( + "Failed to open Zarr group at {}: {}", + object.location, e + )) + })?; + + match Self::find_partitioned_files(&group).await { + Some(partition_groups) => { + if partition_groups.is_empty() { + // No sub-groups found, treat as single group + let statistics = pushdown_statistics::generate_statistics_from_zarr_group( + table_schema, + &self.zarr_pushdown_statistics, + Arc::new(group), + ) + .await + .unwrap_or(Statistics::new_unknown(table_schema)); + + Ok(FileGroup::new(vec![ + PartitionedFile::new(group_path.as_zarr_json_path(), 0) + .with_statistics(Arc::new(statistics)), + ])) + } else { + // Multiple sub-groups found, treat each as a partition + let mut files = Vec::new(); + for group in partition_groups { + let group = Arc::new(group); + let partition_path = group.path().to_string(); + let statistics = pushdown_statistics::generate_statistics_from_zarr_group( + table_schema, + &self.zarr_pushdown_statistics, + group.clone(), + ) + .await + .unwrap_or(Statistics::new_unknown(table_schema)); + + // Create a PartitionedFile for the sub-group + let partitioned_file = + PartitionedFile::new(format!("{}/zarr.json", partition_path), 0) + .with_statistics(Arc::new(statistics)); + files.push(partitioned_file); + } + Ok(FileGroup::new(files)) + } + } + None => Ok(FileGroup::new(vec![])), + } + } + + async fn find_partitioned_files( + group: &Group, + ) -> Option>> { + match group.async_child_groups().await { + Ok(children) => { + // Find recursively all child groups that contain groups aswell + let mut result = Vec::new(); + for child in children { + if let Some(mut sub_children) = + Box::pin(Self::find_partitioned_files(&child)).await + { + result.append(&mut sub_children); + } else { + result.push(child); + } + } + Some(result) + } + Err(_) => None, + } + } +} + +#[cfg(test)] +mod tests { + use datafusion::datasource::listing::ListingTableUrl; + use futures::StreamExt; + + #[test] + fn test_name() { + let test_path = object_store::path::Path::parse("bucket/key.zarr").unwrap(); + + println!("Path: {:?}", test_path); + } + + #[tokio::test] + async fn test_listing_table() { + let object_store = + object_store::local::LocalFileSystem::new_with_prefix("./test_files").unwrap(); + + let url = ListingTableUrl::try_new( + "file:///".try_into().unwrap(), + Some(glob::Pattern::new("gridded-example.zarr").unwrap()), + ) + .unwrap(); + println!("URL: {:?}", url); + + let session = datafusion::execution::context::SessionContext::new(); + let state = session.state(); + let mut files = url + .list_all_files(&state, &object_store, "zarr.json") + .await + .unwrap(); + + while let Some(file) = files.next().await { + println!("File: {:?}", file); + } + + let url = + ListingTableUrl::try_new("file:///gridded-example.zarr/".try_into().unwrap(), None) + .unwrap(); + println!("URL: {:?}", url); + + let session = datafusion::execution::context::SessionContext::new(); + let state = session.state(); + let mut files = url + .list_all_files(&state, &object_store, "zarr.json") + .await + .unwrap(); + + while let Some(file) = files.next().await { + println!("File: {:?}", file); + } + } +} diff --git a/beacon-formats/src/zarr/opener.rs b/beacon-formats/src/zarr/opener.rs new file mode 100644 index 00000000..22261ac6 --- /dev/null +++ b/beacon-formats/src/zarr/opener.rs @@ -0,0 +1,332 @@ +use std::{collections::HashMap, pin::Pin, sync::Arc}; + +use arrow::{datatypes::SchemaRef, error::ArrowError}; +use beacon_arrow_zarr::{ + array_slice_pushdown::ArraySlicePushDown, reader::AsyncArrowZarrGroupReader, + stream::ArrowZarrStream, +}; +use datafusion::{ + common::pruning::PrunableStatistics, + datasource::{ + listing::PartitionedFile, + physical_plan::{FileMeta, FileOpenFuture, FileOpener, ParquetSource}, + schema_adapter::SchemaAdapter, + }, + physical_optimizer::pruning::{FilePruner, PruningPredicate}, + physical_plan::PhysicalExpr, +}; +use futures::{FutureExt, StreamExt}; +use object_store::ObjectStore; +use parking_lot::Mutex; +use zarrs::group::Group; +use zarrs_object_store::AsyncObjectStore; +use zarrs_storage::AsyncReadableListableStorageTraits; + +use crate::zarr::{ + array_step_span::NumericArrayStepSpan, + expr_util::extract_range_from_physical_filters, + pushdown_statistics::ZarrPushDownStatistics, + source::recursive_groups, + stream_share::{PartitionedZarrStreamShare, ZarrStreamShare}, + util::ZarrPath, +}; + +/// Open a Zarr object and produce a future that yields a stream proxy for its partitions. +/// +/// This implements FileOpener for ZarrFileOpener and performs the following steps: +/// - Converts the provided FileMeta's ObjectMeta into a ZarrPath, returning a DataFusion +/// Execution error if parsing fails. +/// - Uses an internal, shared map of stream shares keyed by object path to obtain an +/// Arc. The map access is protected by a lock; if a share for the +/// path does not exist it is created lazily. This ensures concurrent open requests +/// for the same path share initialization work and the resulting streams. +/// - Builds and returns a boxed asynchronous future. When polled, that future: +/// - Calls get_or_try_init on the ZarrStreamShare to initialize or reuse the +/// partitioned streams for the object. Initialization is performed by +/// fetch_partitioned_streams which may access the object store, apply pruning +/// using pushdown statistics/predicates, and consider array step configuration. +/// - Flattens the resulting partitioned streams into a single stream proxy suitable +/// for DataFusion consumption and returns it. +/// +/// Parameters: +/// - file_meta: Metadata for the file to open; its ObjectMeta is used to construct +/// the ZarrPath and to key the stream-share cache. +/// - _file: A PartitionedFile value that is unused by this opener (present to match +/// the trait signature). +/// +/// Return: +/// - On success: Ok(BoxFuture) that resolves to a stream proxy of the opened file's +/// partitions. +/// - On error: DataFusionError::Execution with a descriptive message (e.g., failed +/// ZarrPath creation or stream initialization). +/// +/// Additional notes: +/// - The method clones necessary handles (object store, schema adapter, predicate, +/// array steps) so the returned future is self-contained and safe to run later. +/// - Debug tracing is emitted for path creation and file open operations to aid +/// diagnosis. +/// - IO and heavy work are deferred until the returned future is awaited/polled. +pub struct ZarrFileOpener { + pub table_schema: SchemaRef, + /// Underlying object store for accessing Zarr data. + pub zarr_object_store: Arc>>, + /// Schema adapter for mapping Zarr schema to Arrow schema. + pub schema_adapter: Arc, + /// Stream partition shares for the Zarr file. + pub stream_partition_shares: + Arc>>>, + /// Array Steps for slicing arrays based on step spans. This is utilized by the pruning predicate pushdown. + pub array_steps: Arc>, + /// Pushdown Zarr Statistics + pub pushdown_zarr_statistics: ZarrPushDownStatistics, + /// Pruning Predicate + pub predicate: Option>, +} + +async fn fetch_partitioned_streams( + zarr_path: ZarrPath, + zarr_object_store: Arc>>, + array_steps: Arc>, + pruning_predicate: Option>, + table_schema_adapter: Arc, +) -> datafusion::error::Result> { + let base_group = Group::async_open( + zarr_object_store.clone() as Arc, + &zarr_path.as_zarr_path(), + ) + .await + .map_err(|e| datafusion::error::DataFusionError::Execution(e.to_string()))?; + + let mut partitioned_groups = Vec::new(); + recursive_groups(Arc::new(base_group), &mut partitioned_groups).await?; + // Iterate over arrays in the group and open them as partitioned groups + + let mut partitioned_streams = Vec::new(); + for group_partition in partitioned_groups { + let reader = AsyncArrowZarrGroupReader::new(group_partition) + .await + .map_err(|e| datafusion::error::DataFusionError::Execution(e.to_string()))?; + + let file_schema = reader.arrow_schema(); + + let (table_mapper, projection) = table_schema_adapter.map_schema(&file_schema)?; + let partition_file_schema = file_schema + .project(&projection) + .map_err(|e| datafusion::error::DataFusionError::ArrowError(Box::new(e), None))?; + + // Generate ArraySlicePushDown based on the predicate and array steps + let array_slice_pushdowns = if let Some(pruning_predicate) = &pruning_predicate { + array_steps + .values() + .filter_map(|span| generate_numeric_span_slice_pushdown(pruning_predicate, span)) + .map(|pd| (pd.dimension().to_string(), pd)) + .collect() + } else { + HashMap::new() + }; + + let stream = reader + .into_parallel_stream_composer(Some(projection), Some(array_slice_pushdowns)) + .map_err(|e| datafusion::error::DataFusionError::Execution(e.to_string()))?; + + let partition_stream = + PartitionedZarrStreamShare::new(stream, table_mapper, Arc::new(partition_file_schema)); + partitioned_streams.push(partition_stream); + } + + Ok(Arc::from(partitioned_streams)) +} + +async fn flatten_partitioned_streams( + partitioned_streams: Arc<[PartitionedZarrStreamShare]>, +) -> Pin< + Box< + dyn futures::Stream> + + Send, + >, +> { + let mut all_streams = Vec::new(); + for partition in partitioned_streams.iter() { + let stream = flatten_batch_stream(partition.stream_composer.pollable_shared_stream()); + all_streams.push(stream); + } + let combined_stream = futures::stream::iter(all_streams).flatten(); + combined_stream.boxed() +} + +fn flatten_batch_stream( + stream: ArrowZarrStream, +) -> Pin< + Box< + dyn futures::Stream> + + Send, + >, +> { + (stream + .map(|res_nd_batch| { + let nd_batch = res_nd_batch.map_err(|e| { + ArrowError::IoError( + "Error reading NdRecordBatch from Zarr stream".to_string(), + std::io::Error::other(e), + ) + })?; + + let arrow_batch = nd_batch + .to_arrow_record_batch() + .map_err(|e| ArrowError::ExternalError(Box::new(e)))?; + + Ok(arrow_batch) + }) + .boxed()) as _ +} + +impl FileOpener for ZarrFileOpener { + fn open( + &self, + file_meta: FileMeta, + file: PartitionedFile, + ) -> datafusion::error::Result { + if let (Some(statistics), Some(predicate)) = (file.statistics, self.predicate.as_ref()) { + let prunable_stats = + PrunableStatistics::new(vec![statistics.clone()], self.table_schema.clone()); + let pruning_predicate = + PruningPredicate::try_new(predicate.clone(), self.table_schema.clone())?; + + let pruning_result = pruning_predicate.prune(&prunable_stats)?; + + // If the file can be pruned, return an empty stream + if pruning_result.iter().all(|&v| !v) { + let empty_stream = futures::stream::empty().boxed(); + return Ok(Box::pin(async move { Ok(empty_stream) })); + } + } + + let zarr_path = ZarrPath::new_from_object_meta(file.object_meta.clone()).map_err(|e| { + datafusion::error::DataFusionError::Execution(format!( + "Failed to create ZarrPath from ObjectMeta: {}", + e + )) + })?; + + let zarr_object_store = self.zarr_object_store.clone(); + let adapter = self.schema_adapter.clone(); + let stream_partition_shares = self.stream_partition_shares.clone(); + let stream_partition_share = { + let mut stream_partition_share_map = stream_partition_shares.lock(); + let object_path = file_meta.object_meta.location.clone(); + tracing::debug!( + "Getting or creating ZarrStreamShare for path: {}", + object_path + ); + stream_partition_share_map + .entry(object_path) + .or_insert_with(|| Arc::new(ZarrStreamShare::new())) + .clone() + }; + let pruning_predicate = self.predicate.clone(); + let array_steps = self.array_steps.clone(); + + // Check if the pruning predicate references any arrays in the pushdown statistics + let fut = async move { + // Determine based on statistics and predicate which arrays to pushdown + + let partitioned_streams = stream_partition_share + .get_or_try_init(|| async move { + tracing::debug!("Opening file: {:?}", file_meta.object_meta.location); + + fetch_partitioned_streams( + zarr_path, + zarr_object_store, + array_steps, + pruning_predicate, + adapter, + ) + .await + }) + .await? + .clone(); + + let stream_proxy = flatten_partitioned_streams(partitioned_streams).await; + Ok(stream_proxy) + }; + + Ok(fut.boxed()) + } +} + +/// Result of translating the current physical pruning predicate(s) into a +/// concrete range for the provided column span. +/// +/// This value is produced by `extract_range_from_physical_filters` and +/// captures the contiguous interval of physical values (or row positions) +/// that are guaranteed to satisfy the pruning predicate applied to +/// `span.column`. The computed range is intended for chunk/block-level +/// pruning so that downstream I/O or scanning can skip data that falls +/// outside the returned interval. +/// +/// Behavior notes: +/// - When the predicate can be mapped to a simple interval, `r` holds that +/// interval (e.g., an inclusive/exclusive min..max). Consumers can use it +/// to restrict reads to the overlap with a chunk's bounds. +/// - If the predicate cannot be expressed as a single contiguous range, or +/// if it is contradictory, the function will indicate that no safe range +/// can be derived (typically via `None` or an empty/invalid range), and +/// no pruning should be applied. +/// - The pruning predicate is provided as a slice here to allow the same +/// translation logic to accept zero, one, or multiple physical filters; +/// callers intentionally pass a one-element slice when converting a single +/// predicate. +/// +/// Callers should always check the returned value before using it to avoid +/// incorrect data exclusion when translation fails or is inexact. +fn generate_numeric_span_slice_pushdown( + pruning_predicate: &Arc, + span: &NumericArrayStepSpan, +) -> Option { + let r = + extract_range_from_physical_filters(std::slice::from_ref(pruning_predicate), &span.column); + let mut min_step = None; + let mut max_step = None; + if let Some(min) = r.as_ref().and_then(|r| r.as_f64_min()) { + // Calculate start index based on step + let start_index = lower_index(span.start, span.step, min); + min_step = Some(if start_index < 0 { 0 } else { start_index }); + } + if let Some(max) = r.as_ref().and_then(|r| r.as_f64_max()) { + // Calculate end index based on step + let end_index = upper_index(span.start, span.step, max); + max_step = Some(if end_index < 0 { 0 } else { end_index }); + } + + // If min_step > max_step, then swap them + if let (Some(min), Some(max)) = (min_step, max_step) + && min > max + { + min_step = Some(max); + max_step = Some(min); + } + + Some(ArraySlicePushDown::new( + span.dimension.clone(), + min_step.map(|v| v as usize), + max_step.map(|v| v as usize), + )) +} + +fn lower_index(start: f64, step: f64, value: f64) -> i64 { + let raw = (value - start) / step; + if step >= 0.0 { + raw.floor() as i64 // ascending: lower = floor + } else { + raw.ceil() as i64 // descending: lower = ceil + } +} + +fn upper_index(start: f64, step: f64, value: f64) -> i64 { + let raw = (value - start) / step; + if step >= 0.0 { + raw.ceil() as i64 // ascending: upper = ceil + } else { + raw.floor() as i64 // descending: upper = floor + } +} diff --git a/beacon-formats/src/zarr/partition.rs b/beacon-formats/src/zarr/partition.rs new file mode 100644 index 00000000..c5e149e5 --- /dev/null +++ b/beacon-formats/src/zarr/partition.rs @@ -0,0 +1,61 @@ +use std::sync::Arc; + +use arrow::datatypes::SchemaRef; +use datafusion::datasource::physical_plan::FileGroup; +use object_store::{ObjectMeta, ObjectStore}; +use zarrs::group::Group; +use zarrs_object_store::AsyncObjectStore; +use zarrs_storage::AsyncReadableListableStorageTraits; + +use crate::zarr::{pushdown_statistics::ZarrPushDownStatistics, util::ZarrPath}; + +pub(crate) async fn find_partitioned_files( + group: &Group, +) -> Option>> { + match group.async_child_groups().await { + Ok(children) => { + // Find recursively all child groups that contain groups aswell + let mut result = Vec::new(); + for child in children { + if let Some(mut sub_children) = Box::pin(find_partitioned_files(&child)).await { + result.append(&mut sub_children); + } else { + result.push(child); + } + } + Some(result) + } + Err(_) => None, + } +} + +pub(crate) async fn partition_zarr_group( + table_schema: SchemaRef, + object: &ObjectMeta, + object_store: Arc, + pushdown_statistics: &ZarrPushDownStatistics, +) -> datafusion::error::Result { + // Create an Zarr AsyncObjectStore wrapper around the ObjectStore + let zarr_store = Arc::new(AsyncObjectStore::new(object_store.clone())); + + let zarr_group_path = ZarrPath::new_from_object_meta(object.clone()).map_err(|e| { + datafusion::error::DataFusionError::Execution(format!( + "Failed to create ZarrPath from object metadata: {}", + e + )) + })?; + + let group = Group::async_open(zarr_store, &zarr_group_path.as_zarr_path()) + .await + .map_err(|e| { + datafusion::error::DataFusionError::Execution(format!( + "Failed to open Zarr group at path '{}': {}", + zarr_group_path.as_zarr_path(), + e + )) + })?; + + + + todo!() +} diff --git a/beacon-formats/src/zarr/pushdown_statistics.rs b/beacon-formats/src/zarr/pushdown_statistics.rs new file mode 100644 index 00000000..9d932e32 --- /dev/null +++ b/beacon-formats/src/zarr/pushdown_statistics.rs @@ -0,0 +1,219 @@ +use std::sync::Arc; + +use arrow::{ + array::AsArray, + datatypes::{Int8Type, SchemaRef}, +}; +use beacon_arrow_zarr::reader::AsyncArrowZarrGroupReader; +use datafusion::{ + common::{ColumnStatistics, Statistics, stats::Precision}, + scalar::ScalarValue, +}; +use zarrs::group::Group; +use zarrs_storage::AsyncReadableListableStorageTraits; + +#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +pub struct ZarrPushDownStatistics { + pub arrays: Vec, +} + +impl ZarrPushDownStatistics { + pub fn has_array(&self, array_name: &str) -> bool { + self.arrays.contains(&array_name.to_string()) + } +} + +pub async fn generate_statistics_from_zarr_group( + schema: &SchemaRef, + pushdown: &ZarrPushDownStatistics, + zarr_group: Arc>, +) -> Option { + let zarr_reader = AsyncArrowZarrGroupReader::new(zarr_group.clone()) + .await + .ok()?; + let mut statistics = Statistics::default(); + for field in schema.fields() { + if pushdown.has_array(field.name()) { + // Read the full array + if let Ok(Some(array)) = zarr_reader.read_array_full(field.name()).await { + // Compute statistics for the array + if let Some(array_stats) = compute_statistics_for_array(&array) { + statistics = statistics.add_column_statistics(array_stats.clone()); + } else { + statistics = statistics.add_column_statistics(ColumnStatistics::new_unknown()); + } + } else { + statistics = statistics.add_column_statistics(ColumnStatistics::new_unknown()); + } + } else { + statistics = statistics.add_column_statistics(ColumnStatistics::new_unknown()); + } + } + + Some(statistics) +} + +fn compute_statistics_for_array(array: &arrow::array::ArrayRef) -> Option { + let data_type = array.data_type().clone(); + let null_count = array.null_count(); + let mut statistics = ColumnStatistics { + null_count: Precision::Exact(null_count), + distinct_count: Precision::Absent, + max_value: Precision::Absent, + min_value: Precision::Absent, + sum_value: Precision::Absent, + }; + + match data_type { + arrow::datatypes::DataType::Boolean => { + let boolean_array = array.as_boolean(); + arrow::compute::min_boolean(boolean_array).map(|min| { + statistics.min_value = Precision::Exact(ScalarValue::Boolean(Some(min))) + }); + arrow::compute::max_boolean(boolean_array).map(|max| { + statistics.max_value = Precision::Exact(ScalarValue::Boolean(Some(max))) + }); + } + arrow::datatypes::DataType::Int8 => { + let int_array = array.as_primitive::(); + arrow::compute::min(int_array) + .map(|min| statistics.min_value = Precision::Exact(ScalarValue::Int8(Some(min)))); + arrow::compute::max(int_array) + .map(|max| statistics.max_value = Precision::Exact(ScalarValue::Int8(Some(max)))); + } + arrow::datatypes::DataType::Int16 => { + let int_array = array.as_primitive::(); + arrow::compute::min(int_array) + .map(|min| statistics.min_value = Precision::Exact(ScalarValue::Int16(Some(min)))); + arrow::compute::max(int_array) + .map(|max| statistics.max_value = Precision::Exact(ScalarValue::Int16(Some(max)))); + } + arrow::datatypes::DataType::Int32 => { + let int_array = array.as_primitive::(); + arrow::compute::min(int_array) + .map(|min| statistics.min_value = Precision::Exact(ScalarValue::Int32(Some(min)))); + arrow::compute::max(int_array) + .map(|max| statistics.max_value = Precision::Exact(ScalarValue::Int32(Some(max)))); + } + arrow::datatypes::DataType::Int64 => { + let int_array = array.as_primitive::(); + arrow::compute::min(int_array) + .map(|min| statistics.min_value = Precision::Exact(ScalarValue::Int64(Some(min)))); + arrow::compute::max(int_array) + .map(|max| statistics.max_value = Precision::Exact(ScalarValue::Int64(Some(max)))); + } + arrow::datatypes::DataType::UInt8 => { + let int_array = array.as_primitive::(); + arrow::compute::min(int_array) + .map(|min| statistics.min_value = Precision::Exact(ScalarValue::UInt8(Some(min)))); + arrow::compute::max(int_array) + .map(|max| statistics.max_value = Precision::Exact(ScalarValue::UInt8(Some(max)))); + } + arrow::datatypes::DataType::UInt16 => { + let int_array = array.as_primitive::(); + arrow::compute::min(int_array) + .map(|min| statistics.min_value = Precision::Exact(ScalarValue::UInt16(Some(min)))); + arrow::compute::max(int_array) + .map(|max| statistics.max_value = Precision::Exact(ScalarValue::UInt16(Some(max)))); + } + arrow::datatypes::DataType::UInt32 => { + let int_array = array.as_primitive::(); + arrow::compute::min(int_array) + .map(|min| statistics.min_value = Precision::Exact(ScalarValue::UInt32(Some(min)))); + arrow::compute::max(int_array) + .map(|max| statistics.max_value = Precision::Exact(ScalarValue::UInt32(Some(max)))); + } + arrow::datatypes::DataType::UInt64 => { + let int_array = array.as_primitive::(); + arrow::compute::min(int_array) + .map(|min| statistics.min_value = Precision::Exact(ScalarValue::UInt64(Some(min)))); + arrow::compute::max(int_array) + .map(|max| statistics.max_value = Precision::Exact(ScalarValue::UInt64(Some(max)))); + } + arrow::datatypes::DataType::Float32 => { + let float_array = array.as_primitive::(); + arrow::compute::min(float_array).map(|min| { + statistics.min_value = Precision::Exact(ScalarValue::Float32(Some(min))) + }); + arrow::compute::max(float_array).map(|max| { + statistics.max_value = Precision::Exact(ScalarValue::Float32(Some(max))) + }); + } + arrow::datatypes::DataType::Float64 => { + let float_array = array.as_primitive::(); + arrow::compute::min(float_array).map(|min| { + statistics.min_value = Precision::Exact(ScalarValue::Float64(Some(min))) + }); + arrow::compute::max(float_array).map(|max| { + statistics.max_value = Precision::Exact(ScalarValue::Float64(Some(max))) + }); + } + arrow::datatypes::DataType::Timestamp(time_unit, _) => match time_unit { + arrow::datatypes::TimeUnit::Second => { + let ts_array = array.as_primitive::(); + arrow::compute::min(ts_array).map(|min| { + statistics.min_value = + Precision::Exact(ScalarValue::TimestampSecond(Some(min), None)) + }); + arrow::compute::max(ts_array).map(|max| { + statistics.max_value = + Precision::Exact(ScalarValue::TimestampSecond(Some(max), None)) + }); + } + arrow::datatypes::TimeUnit::Millisecond => { + let ts_array = array.as_primitive::(); + arrow::compute::min(ts_array).map(|min| { + statistics.min_value = + Precision::Exact(ScalarValue::TimestampMillisecond(Some(min), None)) + }); + arrow::compute::max(ts_array).map(|max| { + statistics.max_value = + Precision::Exact(ScalarValue::TimestampMillisecond(Some(max), None)) + }); + } + arrow::datatypes::TimeUnit::Microsecond => { + let ts_array = array.as_primitive::(); + arrow::compute::min(ts_array).map(|min| { + statistics.min_value = + Precision::Exact(ScalarValue::TimestampMicrosecond(Some(min), None)) + }); + arrow::compute::max(ts_array).map(|max| { + statistics.max_value = + Precision::Exact(ScalarValue::TimestampMicrosecond(Some(max), None)) + }); + } + arrow::datatypes::TimeUnit::Nanosecond => { + let ts_array = array.as_primitive::(); + arrow::compute::min(ts_array).map(|min| { + statistics.min_value = + Precision::Exact(ScalarValue::TimestampNanosecond(Some(min), None)) + }); + arrow::compute::max(ts_array).map(|max| { + statistics.max_value = + Precision::Exact(ScalarValue::TimestampNanosecond(Some(max), None)) + }); + } + }, + arrow::datatypes::DataType::Binary => { + let binary_array = array.as_binary::(); + arrow::compute::min_binary(binary_array).map(|min| { + statistics.min_value = Precision::Exact(ScalarValue::Binary(Some(min.to_vec()))) + }); + arrow::compute::max_binary(binary_array).map(|max| { + statistics.max_value = Precision::Exact(ScalarValue::Binary(Some(max.to_vec()))) + }); + } + arrow::datatypes::DataType::Utf8 => { + let string_array = array.as_string::(); + arrow::compute::min_string(string_array).map(|min| { + statistics.min_value = Precision::Exact(ScalarValue::Utf8(Some(min.to_string()))) + }); + arrow::compute::max_string(string_array).map(|max| { + statistics.max_value = Precision::Exact(ScalarValue::Utf8(Some(max.to_string()))) + }); + } + _ => {} // Unsupported data type for statistics + } + + Some(statistics) +} diff --git a/beacon-formats/src/zarr/source.rs b/beacon-formats/src/zarr/source.rs new file mode 100644 index 00000000..e77ce454 --- /dev/null +++ b/beacon-formats/src/zarr/source.rs @@ -0,0 +1,304 @@ +use std::{collections::HashMap, sync::Arc}; + +use arrow::{ + array::RecordBatch, + datatypes::{Schema, SchemaRef}, + error::ArrowError, +}; +use beacon_arrow_zarr::{ + array_slice_pushdown::ArraySlicePushDown, reader::AsyncArrowZarrGroupReader, +}; +use datafusion::{ + common::Statistics, + config::ConfigOptions, + datasource::{ + listing::PartitionedFile, + physical_plan::{ + FileGroup, FileMeta, FileOpenFuture, FileOpener, FileScanConfig, FileSource, + }, + schema_adapter::{DefaultSchemaAdapterFactory, SchemaAdapter, SchemaAdapterFactory}, + }, + physical_expr::conjunction, + physical_plan::{ + PhysicalExpr, + filter_pushdown::{FilterPushdownPropagation, PushedDown}, + metrics::ExecutionPlanMetricsSet, + }, +}; +use futures::{FutureExt, StreamExt}; +use object_store::{ObjectMeta, ObjectStore}; +use parking_lot::Mutex; +use zarrs::group::Group; +use zarrs_object_store::AsyncObjectStore; +use zarrs_storage::AsyncReadableListableStorageTraits; + +use crate::zarr::{ + array_step_span::NumericArrayStepSpan, + expr_util::extract_range_from_physical_filters, + opener::ZarrFileOpener, + pushdown_statistics::ZarrPushDownStatistics, + stream_share::ZarrStreamShare, + util::{ZarrPath, path_parent}, +}; + +pub async fn fetch_schema( + object_store: Arc, + zarr_path: &ZarrPath, +) -> datafusion::error::Result { + let zarr_store = Arc::new(AsyncObjectStore::new(object_store)) + as Arc; + + let group = Group::async_open(zarr_store.clone(), &zarr_path.as_zarr_path()) + .await + .map_err(|e| datafusion::error::DataFusionError::Execution(e.to_string()))?; + let mut zarr_groups = Vec::new(); + recursive_groups(Arc::new(group), &mut zarr_groups).await?; + + // For each zarr group, fetch its schema and merge them. + let mut schemas = Vec::new(); + for zarr_group in zarr_groups { + let reader = AsyncArrowZarrGroupReader::new(zarr_group.clone()) + .await + .map_err(|e| datafusion::error::DataFusionError::Execution(e.to_string()))?; + let schema = reader.arrow_schema(); + schemas.push(schema); + } + + let super_schema = beacon_common::super_typing::super_type_schema(&schemas).map_err(|e| { + datafusion::error::DataFusionError::Execution(format!( + "Failed to compute super schema for Zarr groups at {}: {}", + zarr_path.as_zarr_path(), + e + )) + })?; + + Ok(Arc::new(super_schema)) +} + +pub async fn recursive_groups( + top_level_group: Arc>, + zarr_groups: &mut Vec>>, +) -> datafusion::error::Result<()> { + let child_groups = top_level_group.async_child_groups().await.map_err(|e| { + datafusion::error::DataFusionError::Execution(format!("Failed to list child groups: {}", e)) + })?; + + if child_groups.is_empty() { + // Add the top-level group itself if it has no child groups. + zarr_groups.push(top_level_group.clone()); + } else { + for child_group in child_groups { + Box::pin(recursive_groups(Arc::new(child_group), zarr_groups)).await?; + } + } + return Ok(()); +} + +#[derive(Default, Clone)] +pub struct ZarrSource { + /// Optional schema adapter factory. + schema_adapter_factory: Option>, + /// Optional schema override. + override_schema: Option, + /// Optional column projection. + projection: Option>, + /// Execution plan metrics. + execution_plan_metrics: ExecutionPlanMetricsSet, + /// Projected statistics. + projected_statistics: Option, + + /// Stream Partition Share + stream_partition_shares: Arc>>>, + /// Array Steps for slicing arrays based on step spans. This is utilized by the pruning predicate pushdown. + array_steps: Arc>, + /// Zarr Statistics + pushdown_zarr_statistics: ZarrPushDownStatistics, + /// Pruning Predicate + predicate: Option>, +} + +impl ZarrSource { + pub fn with_array_steps(mut self, array_steps: HashMap) -> Self { + self.array_steps = Arc::new(array_steps); + self + } + + pub fn with_pushdown_zarr_statistics( + mut self, + pushdown_zarr_statistics: ZarrPushDownStatistics, + ) -> Self { + self.pushdown_zarr_statistics = pushdown_zarr_statistics; + self + } +} + +impl FileSource for ZarrSource { + /// Creates a file opener for the given object store and scan config. + fn create_file_opener( + &self, + object_store: Arc, + base_config: &FileScanConfig, + _partition: usize, + ) -> Arc { + let table_schema = self + .override_schema + .clone() + .unwrap_or_else(|| base_config.file_schema.clone()); + let projected_schema = base_config.projected_schema(); + let schema_adapter_factory = self + .schema_adapter_factory + .clone() + .unwrap_or_else(|| Arc::new(DefaultSchemaAdapterFactory)); + let schema_adapter = schema_adapter_factory.create(projected_schema, table_schema.clone()); + let arc_schema_adapter: Arc = Arc::from(schema_adapter); + + Arc::new(ZarrFileOpener { + table_schema, + zarr_object_store: Arc::new(AsyncObjectStore::new(object_store)), + schema_adapter: arc_schema_adapter, + stream_partition_shares: self.stream_partition_shares.clone(), + array_steps: self.array_steps.clone(), + predicate: self.predicate.clone(), + pushdown_zarr_statistics: self.pushdown_zarr_statistics.clone(), + }) + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn with_batch_size(&self, _batch_size: usize) -> Arc { + Arc::new(self.clone()) + } + + fn with_schema(&self, schema: SchemaRef) -> Arc { + Arc::new(Self { + override_schema: Some(schema), + projection: self.projection.clone(), + execution_plan_metrics: self.execution_plan_metrics.clone(), + projected_statistics: self.projected_statistics.clone(), + schema_adapter_factory: self.schema_adapter_factory.clone(), + stream_partition_shares: self.stream_partition_shares.clone(), + array_steps: self.array_steps.clone(), + predicate: self.predicate.clone(), + pushdown_zarr_statistics: self.pushdown_zarr_statistics.clone(), + }) + } + + fn with_projection(&self, config: &FileScanConfig) -> Arc { + Arc::new(Self { + override_schema: self.override_schema.clone(), + projection: config.projection.clone(), + execution_plan_metrics: self.execution_plan_metrics.clone(), + projected_statistics: self.projected_statistics.clone(), + schema_adapter_factory: self.schema_adapter_factory.clone(), + stream_partition_shares: self.stream_partition_shares.clone(), + array_steps: self.array_steps.clone(), + predicate: self.predicate.clone(), + pushdown_zarr_statistics: self.pushdown_zarr_statistics.clone(), + }) + } + + fn with_statistics(&self, statistics: Statistics) -> Arc { + Arc::new(Self { + override_schema: self.override_schema.clone(), + projection: self.projection.clone(), + execution_plan_metrics: self.execution_plan_metrics.clone(), + projected_statistics: Some(statistics), + schema_adapter_factory: self.schema_adapter_factory.clone(), + stream_partition_shares: self.stream_partition_shares.clone(), + array_steps: self.array_steps.clone(), + predicate: self.predicate.clone(), + pushdown_zarr_statistics: self.pushdown_zarr_statistics.clone(), + }) + } + + fn metrics(&self) -> &ExecutionPlanMetricsSet { + &self.execution_plan_metrics + } + + fn statistics(&self) -> datafusion::error::Result { + if let Some(statistics) = &self.projected_statistics { + Ok(statistics.clone()) + } else if let Some(schema) = self.override_schema.as_ref() { + Ok(Statistics::new_unknown(schema)) + } else { + Err(datafusion::error::DataFusionError::Execution( + "Schema must be set to compute statistics".to_string(), + )) + } + } + + fn file_type(&self) -> &str { + "zarr" + } + + fn with_schema_adapter_factory( + &self, + factory: Arc, + ) -> datafusion::error::Result> { + Ok(Arc::new(Self { + override_schema: self.override_schema.clone(), + projection: self.projection.clone(), + execution_plan_metrics: self.execution_plan_metrics.clone(), + projected_statistics: self.projected_statistics.clone(), + schema_adapter_factory: Some(factory), + stream_partition_shares: self.stream_partition_shares.clone(), + array_steps: self.array_steps.clone(), + predicate: self.predicate.clone(), + pushdown_zarr_statistics: self.pushdown_zarr_statistics.clone(), + })) + } + + fn schema_adapter_factory(&self) -> Option> { + self.schema_adapter_factory.clone() + } + + fn repartitioned( + &self, + target_partitions: usize, + repartition_file_min_size: usize, + output_ordering: Option, + config: &FileScanConfig, + ) -> datafusion::error::Result> { + // Repartition by duplicating the file groups to reach the target number of partitions. + let file_groups = config.file_groups.clone(); + + let repartitioned: Vec = file_groups + .iter() + .cycle() + .take(target_partitions) + .cloned() + .collect(); + + Ok(config.clone().with_file_groups(repartitioned).into()) + } + + fn try_pushdown_filters( + &self, + filters: Vec>, + _config: &ConfigOptions, + ) -> datafusion::error::Result>> { + let Some(file_schema) = self.override_schema.clone() else { + return Ok(FilterPushdownPropagation::with_parent_pushdown_result( + vec![PushedDown::No; filters.len()], + )); + }; + + let predicate = match self.predicate.clone() { + Some(predicate) => conjunction(std::iter::once(predicate).chain(filters.clone())), + None => conjunction(filters.clone()), + }; + + let source = Self { + predicate: Some(predicate), + ..self.clone() + }; + + Ok(FilterPushdownPropagation::with_parent_pushdown_result(vec![ + PushedDown::No; + filters.len() + ]) + .with_updated_node(Arc::new(source))) + } +} diff --git a/beacon-formats/src/zarr/stream_share.rs b/beacon-formats/src/zarr/stream_share.rs new file mode 100644 index 00000000..4bc7b725 --- /dev/null +++ b/beacon-formats/src/zarr/stream_share.rs @@ -0,0 +1,55 @@ +use std::{future::Future, sync::Arc}; + +use arrow::datatypes::SchemaRef; +use beacon_arrow_zarr::stream::ArrowZarrStreamComposerRef; +use datafusion::datasource::schema_adapter::SchemaMapper; +use futures::FutureExt; + +pub struct PartitionedZarrStreamShare { + pub stream_composer: ArrowZarrStreamComposerRef, + pub table_schema_mapper: Arc, + pub partition_file_schema: SchemaRef, +} + +impl PartitionedZarrStreamShare { + pub fn new( + stream_share: ArrowZarrStreamComposerRef, + table_schema_mapper: Arc, + partition_file_schema: SchemaRef, + ) -> Self { + Self { + stream_composer: stream_share, + table_schema_mapper, + partition_file_schema, + } + } +} + +pub struct ZarrStreamShare { + partitions: tokio::sync::OnceCell>, +} + +impl ZarrStreamShare { + pub fn new() -> Self { + Self { + partitions: tokio::sync::OnceCell::new(), + } + } + + pub fn get_or_try_init(&self, f: F) -> impl Future + where + F: FnOnce() -> Fut, + Fut: Future< + Output = Result, datafusion::error::DataFusionError>, + >, + { + // Clone the Arc inside the OnceCell to return it. + self.partitions.get_or_try_init(f).map(|f| f.cloned()) + } +} + +impl Default for ZarrStreamShare { + fn default() -> Self { + Self::new() + } +} diff --git a/beacon-formats/src/zarr/util.rs b/beacon-formats/src/zarr/util.rs new file mode 100644 index 00000000..d7f4f922 --- /dev/null +++ b/beacon-formats/src/zarr/util.rs @@ -0,0 +1,241 @@ +use std::collections::{HashMap, HashSet}; + +use object_store::ObjectStore; + +#[derive(Clone, Debug)] +pub enum ZarrPath { + ObjectMeta(object_store::ObjectMeta), + // Directory path representing a Zarr group. Inside there should always be a "zarr.json" file. + DirPath(object_store::path::Path), +} + +impl ZarrPath { + pub fn as_zarr_path(&self) -> String { + let directory_path = match self { + ZarrPath::ObjectMeta(meta) => { + // Strip the trailing "zarr.json" to get the group path. + let loc = meta.location.as_ref(); + loc.strip_suffix("zarr.json").unwrap_or(loc).to_string() + } + ZarrPath::DirPath(path) => { + // Check if ends with '/' and return as is or add it. + path.as_ref().to_string() + } + }; + let p_str = directory_path.trim_end_matches('/'); + format!("/{}", p_str) + } + + pub fn as_zarr_json_path(&self) -> String { + match self { + ZarrPath::ObjectMeta(meta) => meta.location.as_ref().to_string(), + ZarrPath::DirPath(path) => { + let zarr_json_path = path.child("zarr.json"); + zarr_json_path.as_ref().to_string() + } + } + } + + pub fn new_from_object_meta(meta: object_store::ObjectMeta) -> Result { + // Ensure this is a zarr.json file? + if !is_zarr_v3_metadata(&meta) { + return Err(format!( + "ObjectMeta at location '{}' is not a Zarr v3 metadata file (zarr.json)", + meta.location.as_ref() + )); + } + Ok(ZarrPath::ObjectMeta(meta)) + } + + pub async fn new_from_dir_path( + object_store: &dyn ObjectStore, + path: object_store::path::Path, + ) -> Result { + // Ensure that path does not end with "zarr.json" and that in the folder underneath there is a zarr.json file. + if path.as_ref().ends_with("zarr.json") { + return Err(format!( + "Provided path '{}' should be a directory path, not a zarr.json file", + path.as_ref() + )); + } + + let zarr_json_path = path.child("zarr.json"); + let exists = object_store.head(&zarr_json_path).await.is_ok(); + if !exists { + return Err(format!( + "No zarr.json file found under directory path '{}'", + path.as_ref() + )); + } + Ok(ZarrPath::DirPath(path)) + } +} + +/// Get parent directory of a Path (S3-style). +/// Example: "a/b/c" -> Some("a/b") +/// "a" -> None +pub fn path_parent(p: &object_store::path::Path) -> Option { + let s = p.to_string(); + if let Some(pos) = s.rfind('/') { + // Parent is everything before the last '/' + let parent_str = &s[..pos]; + Some(object_store::path::Path::from(parent_str)) + } else { + // No '/' in path => parent is root (optional: return None instead) + None + } +} + +/// Check if this ObjectMeta represents a Zarr v3 metadata file ("zarr.json") +pub fn is_zarr_v3_metadata(meta: &object_store::ObjectMeta) -> bool { + // Normalize for safety, S3 paths are UTF-8 so this is fine + let loc = meta.location.to_string().to_lowercase(); + loc.ends_with("/zarr.json") +} + +/// Return only the ObjectMeta entries corresponding to **top-level Zarr groups**. +pub fn top_level_zarr_meta_v3(metas: &[object_store::ObjectMeta]) -> Vec { + // 1. Collect all metas that are zarr.json + record their directories. + let mut dir_to_meta: HashMap = + HashMap::new(); + + for meta in metas { + if is_zarr_v3_metadata(meta) + && let Some(parent) = path_parent(&meta.location) + { + dir_to_meta.insert(parent, meta); + } + } + + // 2. Extract all candidate directories and sort. + let mut candidates: Vec = dir_to_meta.keys().cloned().collect(); + candidates.sort(); + candidates.dedup(); + + // 3. Build a set for quick ancestor checks. + let all_dirs: HashSet = candidates.iter().cloned().collect(); + + // 4. Keep only directories that have no ancestor also in all_dirs. + let mut top_level_dirs: Vec = Vec::new(); + + 'outer: for dir in &candidates { + let mut current = path_parent(dir); + while let Some(ancestor) = current { + if ancestor == *dir { + break; + } + if all_dirs.contains(&ancestor) { + // dir is nested under another zarr.json directory + continue 'outer; + } + current = path_parent(&ancestor); + } + top_level_dirs.push(dir.clone()); + } + + // 5. Return the corresponding ObjectMetas (cloned) + top_level_dirs + .into_iter() + .filter_map(|d| dir_to_meta.get(&d).cloned().cloned()) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + use object_store::{ObjectMeta, path::Path}; + + fn meta(path: &str) -> ObjectMeta { + ObjectMeta { + location: Path::from(path), + last_modified: Default::default(), + size: 0, + e_tag: None, + version: None, + } + } + + fn extract_paths(metas: &[ObjectMeta]) -> Vec { + metas.iter().map(|m| m.location.to_string()).collect() + } + + #[test] + fn test_single_top_level() { + let metas = vec![meta("a/zarr.json")]; + let result = top_level_zarr_meta_v3(&metas); + assert_eq!(extract_paths(&result), vec!["a/zarr.json"]); + } + + #[test] + fn test_nested_group_dropped() { + let metas = vec![ + meta("a/zarr.json"), // top-level + meta("a/b/zarr.json"), // nested -> drop + ]; + let result = top_level_zarr_meta_v3(&metas); + assert_eq!(extract_paths(&result), vec!["a/zarr.json"]); + } + + #[test] + fn test_multiple_top_level() { + let metas = vec![meta("a/zarr.json"), meta("b/zarr.json")]; + let result = top_level_zarr_meta_v3(&metas); + let mut paths = extract_paths(&result); + paths.sort(); + assert_eq!(paths, vec!["a/zarr.json", "b/zarr.json"]); + } + + #[test] + fn test_deep_nesting_dropped() { + let metas = vec![ + meta("a/zarr.json"), + meta("a/b/zarr.json"), + meta("a/b/c/zarr.json"), + ]; + let result = top_level_zarr_meta_v3(&metas); + assert_eq!(extract_paths(&result), vec!["a/zarr.json"]); + } + + #[test] + fn test_sibling_groups() { + let metas = vec![ + meta("root1/zarr.json"), + meta("root2/zarr.json"), + meta("root2/child/zarr.json"), // nested -> drop + ]; + let result = top_level_zarr_meta_v3(&metas); + let mut paths = extract_paths(&result); + paths.sort(); + assert_eq!(paths, vec!["root1/zarr.json", "root2/zarr.json"]); + } + + #[test] + fn test_root_level_key() { + let metas = vec![ + meta("zarr.json"), // top-level at root + meta("a/zarr.json"), // nested under root -> drop + ]; + let result = top_level_zarr_meta_v3(&metas); + assert_eq!(extract_paths(&result), vec!["a/zarr.json"]); + } + + #[test] + fn test_ignore_non_zarr_json() { + let metas = vec![ + meta("a/zarr.json"), // valid + meta("b/not_zarr.json"), + meta("c/zarr.txt"), + meta("d/data.bin"), + ]; + let result = top_level_zarr_meta_v3(&metas); + assert_eq!(extract_paths(&result), vec!["a/zarr.json"]); + } + + #[test] + fn test_empty_input() { + let metas: Vec = vec![]; + let result = top_level_zarr_meta_v3(&metas); + assert!(result.is_empty()); + } +} diff --git a/zarr.ipynb b/zarr.ipynb new file mode 100644 index 00000000..e69de29b