What are the ultimate aims of `arrow-zarr`?

[JackKelly]

Hi! Please forgive my very naive question! But please may I ask: What are the ultimate aims of arrow-zarr?

I ask because I'm toying with the idea of helping to speed up Python's xarray package. Perhaps by writing some of xarray in Rust. The end-goal would be much faster execution on a single machine by optimising queries, using all CPU cores for processing, and overlapping IO with compute. (None of these ideas are original to me, of course! And some of these ideas are already implemented in packages like dask, cubed, zarrs, and IceChunk.)

It occured to me that, perhaps, it'd be possible to use datafusion to do this! Specifically, maybe datafusion could be taught how to process labelled multi-dimensional arrays (with the appropriate custom code, of course). Within datafusion, we'd pass around arrow::Tensor structs. The ultimate ambition might be to re-implement a sizable portion of xarray in datafusion, and then expose an xarray-like API to Python.

And then I remembered that @maximedion2 had mentioned arrow-zarr to me a year ago, back when I was playing with io_uring! (And, in all likelihood, it was @maximedion2's comment that had primed me to think about datafusion for multi-dimensional arrays!).

So, I was wondering: Have you already built what I've been dreaming about?! Is arrow-zarr already an attempt to build something like xarray in datafusion?! (Absolutely no worries either way! I don't want to add pressure!)

[maximedion2]

Hey there @JackKelly!

So in short, my aim is to expose a query engine (SQL plus maybe a few custom things) that would work on top of zarr storage. More below on why I think that could be useful and a few features I'd like to implement, but first some context.

I started working on this a long time ago, with the goal of learning more about Rust and about zarr and really a stretch goal of actually implementing the query engine. It's tangentially related to what I do at work, but not enough that my company would pay me to work on it so I do all the work on my own time. I managed to implement a simple first version of what I wanted, including code to fully decode and read zarr v2/v3 data (simple file storage, not e.g. Icechunk or anything like that, and it's messy!). There were a lot of other things that I wanted to implement but then... I had a baby! With waaaayy less time to work on this outside my work hours, progress stalled a bit, but I'm trying to get back to it now!

(side note: some of what I had initially written was using io-uring to read chunks asynchronously, that was very heavily inspired (some would say copied a little bit!) from the code you wrote in lightpseed-io @JackKelly, hence why I was bugging you with questions!).

I think zarrs has become the "canonical" zarr Rust implementation, so I've decided to completely rewrite everything to work with that. Implementing all of it myself was a fun learning experience, but I should of course use the reference implementation, not some hacky version of it I put together to learn about zarr. That will make the code much cleaner and shorter, and allow the code to work with e.g. Icechunk.

In terms of implementation, this is all based on datafusion and RecordBatches (which is an implementation of the arrow in-memory format, but the "arrow-zarr" name might not be the most descriptive... oh well). I'm following the same overall pattern as with parquet or csv. So one clarification is that after the data is initially ingested, there is no concept of "arrays", it's just columnar data. That makes it easy to fully leverage everything that datafusion comes with, and when I want to implement new functions, that might get fed into logical/execution plan optimizers, I can again re-use a lot of what already exists. But, provided the right data is present, data in columnar format is only a sort away from being array data, so I was thinking that the results of a query could be written back to a zarr store.

A few reasons I thought this could potentially be a powerful tool:

• The fact that zarr data is chunked up means you can pretty much always assume you have enough concurrent streams to saturate any number of cpus and cores, without blowing up your memory usage.
• I think there's a python package, datafusion-ray, that even allows for distributed computing (not sure if it works "out-of-the-box" for custom datafusion stuff though).
• Because each "variable" (array) is stored in its own file, you can do some pretty efficient stuff with filter pushdowns, even with compressed data (you can read just the chunks for the variables in the filter, and skip all the other variables when appropriate).
• Something that I'm particularly interested in (but requires a good amount of work) is that I think the structure of the zarr chunks, say as a 2D or 3D array with coordinates, is pretty much perfect for r-tree (or r*-tree) spatial indexing. So one could optionally get an indexing "hint" from reading the data, and carrying this hint into some geospatial functions could allow for blazing fast geospatial joins against polygons read as wkb objects from a parquet file (I might even try to implement a point-to-point "join nearest", but that might be ambitious).

A couple relevant use cases for the last point would be interpolating an array to different resolutions (point-to-point "join nearest") or intersecting fairly complex wildfire burn scars against extensive weather/vegetation/topography data stores in the zarr format (to provide training "slices" for ML models, for example).

So yeah that's the dream! I'm not sure how much progress I will be able to make and how fast, but I'm giving it a shot! Does that at least partially overlap with what you had in mind?

[JackKelly]

This is super-interesting, thank you so much for the fast and detailed (and well-written) reply! (And congratulations on recently becoming a parent!)

In short: I definitely think your aims overlap with my aims (which is great!). But we might have slightly different users in mind. (Which is also great, because hopefully we can work on complementary parts of the solution!).

I should also mention that, like you, I'm not sure how much time I'll be able to spend on this stuff. We're currently pitching for a large project at work. If we get the project then I'll probably have to hide away for a while! If we don't get the project then I should have more time.

Following your lead, I should also give some context: My background is in building machine learning systems to forecast wind power and solar power. We use lots of gridded numerical weather predictions (NWP) and satellite data. So, I have three main uses for geospatial tools (like xarray):

1. Before ML training: To convert raw NWP datasets and satellite datasets to on-disk data structures that are easier to train ML models from.
2. During ML training: To select "examples" whilst training ML models.
3. At inference time: To pull live NWP & satellite data, and to feed this data into our ML pipeline.

A big problem at the moment is that, because modern GPUs are so fast, it turns out that the bottleneck during ML training is often creating the training examples. Over the years, we've tried all sorts of horrible hacky workarounds (like preparing batches ahead-of-time), which always come with horrible side-effects like massively slowing down the rate at which we can try out new ML experiments. And, even with these work-arounds, we still need to train on costly machines with huge numbers of CPU cores because the data processing code isn't very computationally efficient.

So I'm keen to try to fix these problems "properly" by helping to make sure we can load data from storage as quickly as possible (hence my interest in io_uring), and helping to process this data as quickly and efficiently as possible.

I and my colleagues like the xarray API, and would be keen to keep using something similar to the xarray API.

We often perform fairly simple operations on data. For example, during the "step 1" that I mentioned above (preparing data prior to training), we might:

1. Rechunk the data. (This can be a surprisingly tricky task to optimise)
2. Rescale all the values so they are between 0 and 255, and convert to uint8.
3. Save to Zarr v3 (either on local SSDs, or on cloud object storage).

Ideally, the user would lazily specify a data processing pipeline that could include multiple on-disk datasets, and joins these datasets, does some computation, and writes to one or more on-disk datasets.

A pipeline for step 2 (training the ML model) might look something like this:

1. Lazily open multiple datasets, e.g.: a few different NWP datasets (from different NWP providers), and a satellite dataset, and a solar power dataset.
2. Each training example will have the same duration (e.g. a 24-hour forecast) and the same geospatial extent (e.g. 64 x 64 pixels). To create a single training example, we randomly generate a start date and central geolocation, and then load data from every dataset that starts at the start date, and is centered on that location. So, label-based selection needs to be super-fast.
3. Then we might downsample.
4. Then we normalise the values (so the mean is zero, and the standard deviation is 1 across the dataset).
5. And we need to generate on the order of 1000 examples per second. So, even with 100 CPU cores, each CPU core has to load and prepare an example every 100 ms.

Ideally, this would all be super-efficient and fast. (One problem with today's toolchain is that dask can easily spend 100 miliseconds just coming up with its "plan"; even though we need to load an entire example in 100 ms!)

So, in summary, I think my aim would be to basically build something that looks almost indistiguishable from xarray to Python users (i.e. it'd use an API that looks a lot like xarray's API, and probably returns an xarray.DataArray). But, under the hood, is optimising queries, minimses IO, overlaps CPU-bound tasks with IO-bound tasks, uses all CPU cores efficiently, and does computations as efficiently as possible.

(In fact, instead of writing a complete xarray-like thing, the ideal solution would be to write a minimal amount of Rust code that hooks into xarray. But I'm not entirely sure if that's possible, because subclassing xarray's ChunkManagerEntrypoint requires us to implement a bunch of methods which take Python functions as arguments... and I'd expect that it'd be super-slow to run Python functions within the Rust backend).

So, after that long rant (!), I guess I should finish up with some concrete questions!

1. Do you think datafusion work for the use-cases I've sketched out above?
2. Is datafusion overkill for my use-cases?! Would it actually be easier to build something from scratch?
3. Are you entirely focused on SQL queries? Or could you imagine users expressing their queries using something like an xarray API?
4. My gut feeling is that it'd be nice to use Arrow's Tensor type to pass around multi-dimensional data. Do you know if that's even possible without with datafusion?!

Thanks again! (And, no rush to reply! Sorry this is a long message!)

[maximedion2]

Ah yes, right I remember reading about trying to saturate GPUs, when I looked at your lightspeed-io project. Ultimately the tasks we use zarr for at my job are quite similar in terms of "machinery", even though the use cases are different scientifically speaking. The most complex tasks we deal with involve training ML models for the spread of wildfires, which more or less require the same 3 high level steps you outlined above.

We're not even using GPUs yet, in fact up until very recently we were using columnar storage just to we could leverage existing tools (spark, duckdb, geopandas, etc...) for the data preparation (your "step 1"), even though our end use case requires raster data, and we definitely "hacked" a lot of things to get that to even work properly. I'm a bit more focused on SQL queries because I'm more familiar with that, having worked with SQL query engines for the last few years, but overall an xarray-like interface would definitely be nice too (more on that below).

Regarding a minimal extension to xarray, I'm not very familiar with the inner workings of xarray and dask, but as you said I'm skeptical about anything that needs to interface with python functions -- that would likely act as a bottleneck that would prevent any meaningful performance gains.

So for your questions, I'll start with 4. You asked about using arrow Tensor without datafusion, did you mean with datafusion? That's not quite possible from what I understand. datafusion has multiple "extension points", but not around what is used to pass data around between operations, it's pretty tightly integrated with RecordBatches. However, there are some things that could be used form datafusion, such as the schemas and some of the more high level query optimization (the LogicalPlan enum in datafusion does not depend on a particular implementation of the data I think). The parts that do depend on RecordBatches could act as a template for how to implement the same logic but based on Tensors instead.

So for 1. and 2., well you'd pretty much have to implement some things from scratch, but datafusion could still help quite a bit, either directly through providing some functionalities and/or through providing examples to help with building others from scratch. One could most likely build something that checks most/all boxes you listed above (optimizing queries, minimizing IO, etc...) around the code and concepts from datafusion.

For 3., while I'm focused on SQL queries, if the datafusion machinery was also implemented around arrow Tensors, it would be great to have both an xarray-like interface and a SQL like interface to tap into the same optimized code. I haven't really thought about what a query language for arrays would look like, most of SQL would probably not directly translate from columns to arrays (hence why for now I'm going for a raster -> arrow -> SQL -> arrow -> raster approach), but that could be interesting and very powerful.

But to be completely honest, I haven't made it far enough into datafusion code yet to provide anything other than high level observations. I'm hoping to get started on implementing custom functions within a month or two, that's where all the fun stuff is, and then I would have a better idea of how all of this could come together. In the end, I think we're trying to build something to solve more or less the same problems as you said.

[maximedion2]

So I looked a bit more into this, when I was working more actively I mostly focused on all the zarr logic (which I will now replace with zarrs) and @tshauck helped be hook that up to datafusion, so I should dig a bit deeper into datafusion anyway.

First, for lazy, CPU efficient, interleaved IO and compute calculations, as I said this is related to the ExecutionPlan trait, specifically the execute method. The RecordBatchStream trait is just a wrapper around a Stream of RecordBatch (plus it adds the concept of a schema), that's the bit that would have to be completely reworked. The trait is implemented on <...>Exec method, for example here it's implemented for a simple limit operation. Basically the struct has a pointer to an input execution plan, which itself exposes a stream that can be polled, and it applies some operation to the record batches it gets and returns them into a stream, which gets fed into the next execution plan, etc...

With the way Rust futures and streams work, my understanding is that nothing is triggered until the top execution plan polls its input, which triggers the whole sequence of operations, so that's your fully lazy evaluation. Also, the "lowest" execution plan, the one used in the FileOpener, doesn't take a stream as an input obviously, it gets the data from the files/zarr stores, and as the data becomes available, it's passed along to the other execution plans. The data is read asynchronously, and the execution plan yields while waiting for IO, so that effectively interleaves IO and compute.

Note that within that very first execution plan, for zarr storage, there's already some compute to decode the chunks, so you can potentially do some interleaving there too, so that chunk n can be decoded while waiting for chunk n+1 to be read. That's what I tried to do, and why I was asking you about interleaving IO and compute, months ago. But, this could potentially be handled differently, with a "read data" execution plan that would feed into a "decode data" execution plan. I might attempt to do that in my refactor actually.

To make use of all the available CPUs, I know that datafusion makes use of partitions to spawn completely independent streams (that is, until they are no longer independent, e.g. you need an average over all the data), and of course as record batches make it along the overall tree of operations multiple CPUs can be used, but other than that I'm not completely sure what's the best way to just keep all the CPUs busy. Obviously with 100s or 1000s of zarr chunks there's potential for keeping tons of CPUs busy, but for some operations, it might be more efficient to accumulate several batches of data together before running the operation, so some "manual" optimization would have to go into this.

Speaking of optimization, in general it seems it's done through implementing the TreeNode trait for logical and execution plans. I'm honestly not too sure how that works, it'll be interesting to look into, but I think that's how the overall tree of operations ends up being re-worked to optimize how everything runs, to that's where all the query optimization logic would go (and in the definition of the logical and execution plans).

Lastly, that execution plan at the very bottom of the tree, the one used in the FileOpener, can take in filters that get pushed down from other execution plans, so that the filters can directly be applied when reading in the data (and some data can be skipped completely). I'm pretty sure xarray already does something like that, but obviously anything that is performance focused will need to skip reading data as much as possible. That step would be particularly important for making your step 2 above as fast as possible.

So basically, you'd need a TensorExecutionPlan trait that deals with Streams of arrow::Tensors, and then you'd implement executors for each tensor based operations. Then you'd need to implement tree traversal for those executors so that the tree can be optimized (you'd also need to deal with a new LogicalPlan implementation I think). I think things would mostly exactly mimic how it works for RecordBatches, but operations would act on tensors.

[JackKelly]

Thank you so much for all the thought and research that you've put into this thread!

You asked about using arrow Tensor without datafusion, did you mean with datafusion?

Oops🤦! Good catch! You're right: I meant "with datafusion". (I've crossed out the "without" in my original post!)

So basically, you'd need a TensorExecutionPlan trait that deals with Streams of arrow::Tensors, and then you'd implement executors for each tensor based operations. Then you'd need to implement tree traversal for those executors so that the tree can be optimized (you'd also need to deal with a new LogicalPlan implementation I think). I think things would mostly exactly mimic how it works for RecordBatches, but operations would act on tensors.

Just to check I understand correctly: Is your suggestion that, if I implemented this stuff, then I could still use datafusion? Or is your proposal that I borrow the core ideas from datafusion, but implement something from scratch?

BTW, a quick update from my end: If I can get time to work on this stuff, then I may start with a smaller project than a complete xarray-compute-backend... I may start with a Rust-based version of rechunker, with a focus on high-performance on a single machine. This idea came from @TomNicholas, during a discussion on the Pangeo forum. I probably won't use datafusion for the first attempt at this Rust rechunker. My understanding is that a fast rechunker would be a useful contribution in its own right, and a good starting point for a more general-purpose xarray-backend-like thingie (if I decide to follow the rabit hole all the way down).

Finally, to correct a mistake I made above: It turns out that it is possible to plug into xarray's backend API from Rust, without having to call Python for every chunk. The Python function calls can just be called once per xarray task, rather than once per chunk (as pointed out by Tom here).

[maximedion2]

I don't think you can use datafusion to implement it. You could (heavily) borrow from datafusion, but you couldn't do it using the "extension" points it exposes, you'd have to re-implement a lot of it for tensors.

Sounds good, working on a smaller scale project to get started is a good idea. On my end, I think something I'm trying to get out of my project, if I can work on it consistently, is to apply geospatial functions to data I get from zarr storage. Someone in that thread you mentioned was talking about GeoRust, which I plan to use if/when I get to that part, and being able to leverage the cloud-friendly format of zarr in combination with efficient geospatial operations is part of what I'm trying to achieve here.

Basically when you're dealing with catastrophic weather events, like wildfires as I mentioned, but also hurricanes, floods, storms, typically the "events" themselves are provided in vector format, not raster format, so I'm trying to come up with something that could help with mapping polygons to existing raster data. There's other stuff too, like reprojecting the data and various forms of aggregation and interpolation, that could be useful. Ultimately I'm hoping this could provide a lot of flexible operations that could be applied to zarr data, but I'll start small too probably!

[JackKelly]

Cool, sounds good! Let's definitely keep in touch. It feels like we definitely have similar interests and aims. And thanks again for doing so much research into whether datafusion could be used for processing multi-dimensional arrays.

[alxmrs]

Max! It looks like we're independently building the same or a very similar idea! Here's my experiment: https://github.com/alxmrs/xarray-sql

Note in particular this issue and the attempt at an implementation (current WIP):

• Zarrquet: Create a referenced based filesystem to read virtual parquet from Zarr. alxmrs/xarray-sql#4
• Zarrquet implementation: Adding Zarrs-based TableProvider with bindings in Python. alxmrs/xarray-sql#69

This is very much aligned with what Jack described at the top of this thread.

Would you like to join forces on this effort? I think a SQL interface on top of Zarr would democratize scientific datasets. I know that these types of operations are tractable (at least, from the example of xvec), and I believe a tabular/SQL interface for array data would bring more people into this space.

[kylebarron]

• Something that I'm particularly interested in (but requires a good amount of work) is that I think the structure of the zarr chunks, say as a 2D or 3D array with coordinates, is pretty much perfect for r-tree (or r*-tree) spatial indexing. So one could optionally get an indexing "hint" from reading the data, and carrying this hint into some geospatial functions could allow for blazing fast geospatial joins against polygons read as wkb objects from a parquet file (I might even try to implement a point-to-point "join nearest", but that might be ambitious).

And FYI I'm working on general vector geospatial support for DataFusion in https://github.com/geoarrow/geoarrow-rs/tree/main/rust/geodatafusion

[alxmrs]

Jack: your use case is exactly the one I was targeting when I started experimenting with SQL, Xarray and Zarr. Before that, I had built a CLI tool in weather-tools (weather-mv) to load xarray-backed rasters into BigQuery. I did it because one of our client teams had an excellent interface for making ML pipelines: my colleague Mohamet would join tabular and "pivoted" raster data together as tables in BQ, which he used to rapidly iterate on which columns/data_vars to include in the featurization pipeline. I agree with Max that a lot of similar geo operations are actually just SQL joins, and that a compelling ML experimentation interface could be built mapping tensor -> table -> SQL query -> table -> tensor. The SQL in the middle is a killer "logical" layer to express these concepts, and with Arrow and Zarrs (and rust-back KDTree indexes, etc.), I think there is good evidence that we can get there in terms of "physical" performance.

My bet is that a SQL interface would bring more people into this ecosystem and toolsets rather than "compete" with things like Xarray.

One way we could weave these tools together initially is for xarray-sql to call into this library to performantly load data from Zarr. From there, one could join Xarray, Dataframes, and pure Zarr stores in the same compute context.

[alxmrs]

The flight that I'm on is about to take off, I'll read these sources on the plane. I'd love to set up a meeting to coordinate development along the idea! I'd love to make it happen! If you're down (or anyone else in the thread), I'd be happy to set up a call over email.

[maximedion2]

Max! It looks like we're independently building the same or a very similar idea! Here's my experiment: https://github.com/alxmrs/xarray-sql

Note in particular this issue and the attempt at an implementation (current WIP):

* [Zarrquet: Create a referenced based filesystem to read virtual parquet from Zarr. alxmrs/xarray-sql#4](https://github.com/alxmrs/xarray-sql/issues/4)

* [Zarrquet implementation: Adding Zarrs-based TableProvider with bindings in Python. alxmrs/xarray-sql#69](https://github.com/alxmrs/xarray-sql/pull/69)

This is very much aligned with what Jack described at the top of this thread.

Would you like to join forces on this effort? I think a SQL interface on top of Zarr would democratize scientific datasets. I know that these types of operations are tractable (at least, from the example of xvec), and I believe a tabular/SQL interface for array data would bring more people into this space.

Hey there!

Yes long story short I do think adding SQL support to Zarr would be great! I think we've collectively listed several reasons to be excited about this!

A few high level remarks on this project:

• The main branch is outdated, I'm working on a side branch where I re-wrote everything in a much cleaner way using zarrs. The implementation on main was me doing everything "zarr related" myself, reading the raw data, handling compression, etc..., which was a cool learning project, but I never intended to really build the library like that, once a proper canonical Rust implementation of the Zarr format came out I switched to that.
• I'm working on this on my own time, can't quite justify doing this for my job (yet), and I have way less time outside of work then I would like (as I told Jack above, the main reason I don't have a lot of time right now just turned 1 year old, is working on her third tooth and is just starting to semi-consistently sleep through the night!). I'm still making some good progress though, just not as much as I'd like!
• As you can see, everything is based on datafusion, this whole library is basically planned to be an extension to datafusion. I'm currently working on ingesting data from zarr/icechunk relatively efficiently, but I also plan on implementing some geospatial functions as well as the option to output the query results to zarr directly, but I still have to look into a lot of datafusionconcepts to get there.

Looking at your project, do I correctly understand that you've recently started considering using datafusion as well for your library?

[maximedion2]

• Something that I'm particularly interested in (but requires a good amount of work) is that I think the structure of the zarr chunks, say as a 2D or 3D array with coordinates, is pretty much perfect for r-tree (or r*-tree) spatial indexing. So one could optionally get an indexing "hint" from reading the data, and carrying this hint into some geospatial functions could allow for blazing fast geospatial joins against polygons read as wkb objects from a parquet file (I might even try to implement a point-to-point "join nearest", but that might be ambitious).

And FYI I'm working on general vector geospatial support for DataFusion in https://github.com/geoarrow/geoarrow-rs/tree/main/rust/geodatafusion

Ah very interesting! Are you planning for this to be a standalone library that extends datafusion with geospatial functions? Are you also considering adding support of geoparquet and geometry objects in datafusion (as opposed to wkt/wkb representations of the geometry from "regular" parquet)?

[maximedion2]

The flight that I'm on is about to take off, I'll read these sources on the plane. I'd love to set up a meeting to coordinate development along the idea! I'd love to make it happen! If you're down (or anyone else in the thread), I'd be happy to set up a call over email.

Sure! I'd like to at least understand what you're planning to do, see if we have complimentary approaches, if we can combine some things, etc... I can also run you through what I've done so far and how I was planning on moving forward with next steps (as I said, I still have some stuff to learn re: datafusion, but it seems to be quite powerful and has many extension points).

[kylebarron]

Ah very interesting! Are you planning for this to be a standalone library that extends datafusion with geospatial functions? Are you also considering adding support of geoparquet and geometry objects in datafusion (as opposed to wkt/wkb representations of the geometry from "regular" parquet)?

Yes, it'll be a standalone library that defines user-defined functions and table providers for DataFusion. So an end product that wants to work with Zarr and geospatial data could register both extensions.

I've recently been working on support for FlatGeobuf and GeoParquet table providers.

Geometries are represented in GeoArrow format. So that means they can be represented in "serialized" format as WKT strings or WKB buffers, or they can be represented in an efficient "native" format that provides constant-time access to any coordinate. The FlatGeobuf and GeoParquet table providers above load data into this GeoArrow format.