This post gives detailed background to my PyData Global talk, “GPU-Accelerated Zarr” (slides, video). It deliberately gets into the weeds, but I will try to provide some background for people who are new to Zarr, GPUs, or both.

The first takeaway is that zarr-python natively supports NVIDIA GPUs. With a one-line zarr.config.enable_gpu() you can configure zarr to return CuPy arrays, which reside on your GPU:

>>> import zarr
>>> zarr.config.enable_gpu()
>>> z = zarr.open_array("path/to/store.zarr", mode="r")
>>> type(z[:])
cupy.ndarray

The second takeaway, and the main focus of this post, is that that simple one-liner leaves performance on the table. It depends a bit on your workload, but I’d claim that Zarr’s data loading pipeline shouldn’t ever be the bottleneck. Achieving maximum throughput today requires some care to ensure that the system’s resources are used efficiently. I’m hopeful that we can improve the libraries to do the right thing in more situations.

This post pairs nicely with Earthmover’s I/O-Maxing Tensors in the Cloud post, which showed that network and object storage service (e.g. S3) also shouldn’t be a bottleneck in most workloads. Ideally, your actual computation is where the majority of time is spent, and the I/O pipeline just gets out of your way.

Some background

I imagine that some people reading this have experience with Zarr but not GPUs, or vice versa. Feel free to skip the sections you’re familiar with, and meet up with us at the Speed of Light section.

Zarr Background for GPU People

Zarr is many things, but today we’ll focus on Zarr as the storage format for n-dimensional arrays. Instead of tabular data, which you might store in a columnar format like Apache Parquet, you’re working with data that fits things like xarray’s data model: everything is an n-dimensional array with metadata. For example, 3-d array measuring forecasts for a temperature field with dimensions (x, y, time).

Zarr is commonly used in many domains including microscopy, genomics, remote sensing, and climate / weather modeling. It works well with both local file systems and remote cloud object storage. High-level libraries like xarray can use zarr as a storage format:

# https://tutorial.xarray.dev/intermediate/remote_data/cmip6-cloud.html
>>> ds = xr.open_zarr(
...     "gs://cmip6/CMIP6/ScenarioMIP/NOAA-GFDL/...",
...     consolidated=True,
... )
>>> zos_2015jan = ds.zos.sel(time="2015-01-16")
>>> zos_2100dec = ds.zos.sel(time="2100-12-16")
>>> sealevelchange = zos_2100dec - zos_2015jan
>>> sealevelchange.plot.imshow()

xarray knows how to translate the high-level slicing like time="2015-01-16" to the lower level slicing of a Zarr array, and Zarr knows how to translate positional slices in the large n-dimensional array to files / objects in storage. This diagram shows the structure of a Zarr store:

The large logical array is split into one or more chunks along one or more dimensions. The chunks are then compressed and stored to disk, which lowers storage costs and can improve read and write performance (it might be faster to read fewer bytes, even if you have to spend time decompressing them).

Zarr’s sharding codec is especially important for GPUs. This makes it possible to store many chunks in the same file (a file on disk, or an object in object storage). We call the collection of chunks a shard, and the shard is what’s actually written to disk.

Multiple chunks are (independently) compressed, concatenated, and stored into the same file / object. We’ll discuss this more when we talk about performance, but the key thing sharding provides is amortizing some constant costs (opening a file, checking its length, etc.) over many chunks, which can be operated on in parallel (which is great news for GPUs).

For now, just note that we’ll be dealing with various levels of Zarr’s hierarchy:

• Arrays: the logical n-dimensional array
• Shards: the file on disk / object in object storage, which contains many chunks concatenated together
• Chunks: the smallest unit we can read (since it must be decompressed to interpret the bytes correctly)

GPU Background for Zarr People

GPUs are massively parallel processors: they excel when you can apply the same problem to a big batch of data. This works well for video games, ML / AI workloads, and data science / data analysis applications.

(NVIDIA) GPUs execute “kernels”, which are essentially functions that run on GPU data. Today, we won’t be discussing how to author a compute kernel. We’ll be using existing kernels (from libraries like nvcomp, CuPy, and CCCL). Instead, we’ll be worried about higher-level things like memory allocations, data movement, and concurrency

Many (though not all) GPU architectures have dedicated GPU memory. This is separate from the regular main memory of your machine (you’ll hear the term “device” to refer to GPUs, and “host” to refer to the host operating system / machine, where your program is running).

While device memory tends to be relatively fast compared to host memory (for example, it might have >3.3 TB/s from the GPU’s memory to its compute cores), it’s moving data between host and device memory is relatively slow (perhaps just 128 GB/s over PCIe). It also tends to be relatively small (an NVIDIA H100 has 80-94GB of GPU memory; newer generations have more, but still GPU memory is precious when processing large datasets). All this means we need to be careful with memory, both how we allocate and deallocate memory and how we move data between the host and device.

In GPU programming, keeping the GPU busy is necessary (but not sufficient!) to achieve good performance. We’ll use GPU utilization, the percent of time (over some window) when the GPU was busy executing some kernel, as a rough measure of how well we’re doing.

One way to achieve high GPU utilization is to queue up work for the GPU to do. The GPU is a device, a coprocessor, onto which your host program offloads work. As much as possible, we’ll have our Python program just do orchestration, leaving the heavy computation to the GPU. Doing this well requires your host program to not slow down the (very fast) GPU.

In some sense, you want your Python program to be “ahead” of the GPU. If you wait to submit your next computation until some data is ready on the GPU, or some previous computation is completed, you’ll inevitably have some time gap when your GPU is idle. Sometimes this is inevitable, but with a bit of care we’ll be able to make our Zarr example perform well.

My Cloud Native Geospatial Conference post touched on this under Pipelining. This program waits to schedule the computation until the CPU is done reading the data, and so doesn’t achieve high throughput:

This second program queues up plenty of work to do, and so achieves higher throughput:

For this example, we’ll use a single threaded program with multiple CUDA Streams to achieve good pipelining. CUDA streams are a way to express a sequence (a stream, if you will) of computations that must happen in order. But, crucially, you can have multiple streams active at the same time. This is nice because it frees you from having to worry too much about exactly how to schedule work on the GPU. For example, one stream of computation might heavily use the memory subsystem (to transfer data from the host to device, for example) while another stream might be using the compute cores. But you don’t have to worry about timing things so that the memory-intensive operation runs at the same time as the compute-intensive operation.

In pseudocode:

a0 = read_chunk("path/to/a", stream=stream_a)
b0 = read_chunk("path/to/b", stream=stream_b)

a1 = transform(a0, stream=stream_a)
b1 = transform(b0, stream=stream_b)

read_chunk might exercise the memory system to transfer data from the host to the device, while transform might really hammer the compute cores.

All you need to do is “just” correctly express the relationships between the different parts of your computation (not always easy!). The GPU will take care of running things concurrently where possible.

One subtle point here: these APIs are typically non-blocking in your host Python (or C/C++/whatever) program. read_chunk makes some CUDA API calls internally to kick off the host to device transfer, but it doesn’t wait for that transfer to complete. This is good, since we want our host program to be well ahead of the GPU; we want to go to the next line and feed the GPU more work to do.

If we actually poked the memory address where the data’s supposed to be it might be junk. We just don’t know. If we really need to wait for some data / computation to be completed, we can call stream.synchronize(), which forces the host program to wait until all the computations on that stream are done. But ideally, you don’t need that. For the typical case of launching some CUDA kernel some some data, synchronization is unnecessary. You only need to ensure that the computation happens on the same CUDA stream as the data loading (like in our pseudocode example, launching each transform on the appropriate stream), and you’re good to go.

CUDA streams do take some getting used to. You can make analogies to thread programming and to async / await, but that only gets you so far. At the end of the day they’re an extremely useful tool to have in your toolkit.

Speed of Light

When analyzing performance, it can be helpful to perform a simple “speed-of-light” analysis: given the constraints of my system, what performance (throughput, latency, whatever metric you care about) should I expect to achieve? This can combine abstract things (like a performance model for how your system operates) with practical things (what’s the sequential read throughput of my disk? What’s the clock cycle of my CPU?).

Many Zarr workloads involve (at least) three stages:

1. Reading bytes from storage (local disk or remote object storage). Your disk (for local storage) or NIC / remote storage service (for remote storage) has some throughput, which you should aim to saturate. Which bytes you need to read will be dictated in part by your application. Zarr supports reading subsets of data (with the chunk being the smallest decompressable unit). Ideally, your chunking should align with your access pattern.

2. Decompressing bytes with the Codec Pipeline. Different codecs have different throughput targets, and these can depend heavily on the data, chunk size, and hardware. We’re using the default Zstd codec in this example.

3. Your actual computation. This should ideally be the bottleneck: it’s the whole reason you’re loading all this data after all.

And if you are using a GPU, at some point you need to get the bytes from host to device memory¹.

Finally, you might need to store your result. If your computation reduces the data this might be negligible. But if you’re outputting large n-dimensional arrays this can be as or more expensive than the reading.

In this case, we don’t really care about what the computation is; just something that uses the data and takes a bit of time. We’ll do a bunch of matrix multiplications because they’re pretty computationally expensive and they’re well suited to GPUs.

Notably, we won’t do any kind computation that involves data from multiple shards. They’re completely independent in this example, which makes parallelizing at the shard level much simpler.

Example Workload

This workload operates on a 1-D float32 array with the following properties:

Each chunk is Zstd compressed, and the shards take about 77.5 MB on disk giving a compression ratio of about 5.3.

The fact that the array is 1-D isn’t too relevant here: zarr supports n-dimension arrays with chunking along any dimension. It does ensure that one optimization is always available when decoding bytes, because the chunks are always contiguous subsets of the shards. We’ll talk about this in detail in the Decode bytes section.

Our workload will read the data, transfer it to the GPU (if using the GPU) and perform a bunch of matrix multiplications.

Performance Summary

This example workload has been fine-tuned to make the GPU look good, and I’ve done zero tuning / optimization of the CPU implementation. Any comparisons with CPU libraries are essentially bunk, but it’s a natural question so I’ll report them anyway.

The top level summary will compare three implementations:

1. zarr-python: Uses vanilla zarr-python for I/O and decoding, and NumPy for the computation.
2. zarr-python GPU: Uses zarr-python’s built-in GPU support to return CuPy arrays, so the GPU is used for computation. At the moment, this still uses Numcodecs to decompress the data, which runs on the CPU. After decompression, the data is moved to the GPU for the matrix multiplication.
3. Custom GPU: My custom implementation of I/O and decoding with CuPy for the computation.

You can find the code for these in my CUDA Stream Samples repository.

Please don’t take the absolute numbers, or even the relative numbers too seriously. I’ve spent zero time optimizing the Zarr/NumPy and Zarr/CuPy implementations. The important thing to take away here is that we have plenty of room for improvement. My Custom I/O pipeline just gradually removed bottlenecks as they came up, some of which apply to zarr-python’s CPU implementation as well. Follow https://github.com/zarr-developers/zarr-python/issues/2904 if you’re interested in developments.

The remainder of the post will describe, in some detail, what makes the custom implementation so fast.

Performance optimizations

Once you have the basics down (using the right data structures / algorithm, removing the most egregious overheads), speeding up a problem often involves parallelization. And you very often have multiple levels of parallelization available. Picking the right level is absolutely a skill that requires some general knowledge about performance and specific details for your problem.

In this case, we’ll operate at the shard level. This will be the maximum amount of data we need to hold in memory at any point in time (though the problem is small enough that we can operate on all the shards at the same time).

We’ll use a few techniques to get good performance in our pipeline:

1. No (large) memory allocations on the critical path.

This applies to both host and device memory allocations. We’ll achieve this by preallocating all the arrays we need to process the shard. Whether or not this should be considered cheating or not is a bit debatable and a bit workload dependent. I’d argue that the most advanced, performance-sensitive workloads will process large amounts of data and so can preallocate a pool of buffers and reuse them across their unit of parallelization (shards in our case).

Regardless, if we’re doing large memory allocations after we’ve started processing a shard (either host or device allocations for the final array or for intermediates) then these allocations can quickly become the bottleneck. Pre-allocation (and reuse across shards) is an important optimization if it’s available.

2. Use pinned (page-locked) memory for host buffers

Using pinned memory makes the host to device transfers much faster. More on that later.

3. Use CUDA streams to overlap I/O and Computation

Our workload has a regular pattern of “read, transfer, decode, compute” on each shard. Because these exercise different parts of the GPU (transfer uses the memory subsystem, decode and compute launch kernels that run on the GPU’s cores), we can run them concurrently.

We’ll assign a CUDA stream per shard. We’ll be very careful to avoid stream / device synchronizations so that our host program schedules all the work to be done.

Throughout this, we’ll use nvtx to annotate certain ranges of code. This will make reading the Nsight Systems report easier.

Here’s a screenshot of an nsys profile, with a few important bits highlighted (open the file for a full-sized screenshot):

• Under Processes > Threads > python, you see the traces for our host program, in this case a Python program. This will include our nvtx annotations (read::disk, read::transfer, read::decode, etc.) and calls to the CUDA API (e.g. cudaMemcpyAsync). These calls measure the time spent by the CPU / host program, not the GPU.

• Under Processes > CUDA HW, you’ll see the corresponding traces for GPU operations. This shows CUDA kernels (functions that run on the GPU) in light blue and memory operations (like host to device transfers) in teal.

You can download the full nsight report here and open it locally with NVIDIA Nsight Systems.

This table summarizes roughly where we spend our time on the GPU per shard (very rough, and there’s some variation across shards, especially as we start overlapping operations with CUDA streams).

Raw throughput measures the actual number of bytes processed per time unit, which is the compressed size for reading, transferring, and decoding. “Effective Throughput” uses the uncompressed number of bytes for each stage. After decompression the actual number of bytes processed equals the uncompressed bytes, so Compute’s raw throughput is equal to its effective throughput.

Read bytes

First, we need to load the data. In my example, I’m just using files on a local disk, though you could use remote object storage and still perform well. We’ll parallelize things at the shard level (i.e. we’re assuming that the entirety of the shard fits in GPU memory).

path = array.store_path.store.root / array.store_path.path / key

with open(path, "rb") as f, nvtx.annotate("read::disk"):
    f.readinto(host_buffer)

On my system, it takes about 13.6 ms to read the 77.5 MB, for a throughput of about 5.7 GB/s from disk (the OS probably had at least some of the pages cached). The effective throughput (uncompressed size over duration) is about 30.1 GB/s. I’ll note that I haven’t spent much effort optimizing this section.

Note that we use readinto to read the data from disk directly into the pre-allocated host buffer: we don’t want any (large) memory allocations on the critical path. Also, we’re using pinned memory (AKA page-locked memory) for the host buffers. This prevents the operating system from paging the buffers, which lets the GPU directly access that memory when copying it, no intermediate buffers required.

And it’s worth emphasizing: this I/O is happening on the host Python program, and it is blocking. As we’ll see later, time spent doing stuff in Python is time not spent scheduling work on the GPU. We’ll need to ensure that the GPU is fed sufficient work, so let’s keep our eye on this section.

The profile report for this section is pretty boring:

Note what the GPU is doing right now: nothing! There aren’t any CUDA HW annotations visible above the initial read::disk. At least for the very first shard we read, the GPU is necessarily idle. But as we’ll discuss shortly, subsequent shards are able to overlap disk I/O with CUDA operations.

This screenshot shows the profile for the second shard:

Now the GPU is busy with some other operations (decoding the chunks from the first shard in this case, which are directly above the read::decode happening on the host at that time). This is partly why I didn’t bother with parallelizing the disk I/O: only one thing can be the bottleneck, and right now we’re able to load data from disk quickly enough.

Transfer bytes

After we’ve read the bytes into memory, we schedule the host to device transfer:

with nvtx.annotate("read::transfer"), stream:
    # device_buffer is a pre-allocated cupy.ndarray
    device_buffer.set(
        host_buffer[:-index_offset].view(device_buffer.dtype), stream=stream
    )

This is where our earlier discussion on blocking vs. non-blocking APIs comes in handy. The device_buffer.set call is not blocking, which is why it takes only ~60 μs on the host. It only makes the CUDA API call to set up the transfer and then immediately returns back to the Python program (to close our context managers and then continue to the next line in our program).

The actual memory copy (which is running on the device) takes about 1.5 ms for a throughput of about 52 GB/s (this is still compressed data, so the effective throughput is even higher). Here’s the same profile I showed earlier, but now you’ll understand the context around what happens on the host (the CUDA API call to do something) and device.

I’ve added the orange lines connecting the fast cudaMemcpyAsync on the host to the (not quite as fast) Memcpy HtoD (host to device) running on the device.

And if you look closely, you’ll see that just above that Memcpy HtoD in teal, we’re executing a compute kernel (in light-blue). We’ll get to that in a bit, but this show that we’re overlapping Host-to-Device transfers with compute kernels.

Decode bytes

At this point we have (or will have, eventually) the Zstd compressed bytes in GPU memory. You might think that “decompressing a stream of bytes” doesn’t mesh well with “GPUs as massively parallel processors”. And you’d be (partially) right! We can’t really parallelize decoding within a single chunk, but we can decode all the chunks in a shard in parallel. My colleague Akshay has a nice overview of how the GPU can be used to decode many buffers in parallel.

I have no idea how to implement a Zstd decompressor, but fortunately we don’t have to. The nvCOMP library implements a bunch of GPU-accelerated compression and decompression routines, including Zstd. It provides C, C++, and Python APIs. A quick note: this example is using a custom wrapper around nvcomp’s C API. This works around a couple issues with nvcomp’s Python bindings.

1. At the moment, accessing an attribute on the decompressed array returned by nvcomp causes a “stream synchronization”. This forces essentially blocks the host program from progressing until the GPU has caught up, which we’d like to avoid. We need to issue compute instructions still, and we’d ideally move on to the next shard!
2. We’d like full control over all the memory allocations, including the ability to preallocate the output buffers that the arrays should be decompressed into. This is possible with the C API, but not (yet) the Python API.

My custom wrapper is not at all robust, well designed, etc. It’s just enough to work for this demo. Don’t use it! Use the official Python bindings, and reach out to me or the nvcomp team if you run into any issues. But here’s the basic idea in code:

zstd_codec = ZstdCodec(stream=stream)
# get a list of arrays, each of which is a view into the original device buffer
# `device_buffer` is stream-ordered on `stream`,
# so `device_arrays` are all stream-ordered on `stream`
device_arrays = [
    device_buffer[offset : offset + size] for offset, size in index
]
with nvtx.annotate("read::decode"):
    zstd_codec.decode_batch(device_arrays, out=out_chunks)

# and now `out` is stream-ordered on `stream`

The zstd_codec.decode_batch call takes about 2.4 ms on my machine. Again this just schedules the decompression call.

The actual decompression takes about 25-45 ms, for a throughput of about roughly 1.7 GB/s.

Again, we’ve pre-allocated the out ndarray, however this is not always possible. Zarr allows chunking over arbitrary dimensions, but we’ve assumed that the chunks are contiguous slices of the output array². If your chunks aren’t contiguous slices of the output array, you’ll need to decode into an intermediate buffer and then perform some memory copies into the output buffer.

Anyway, all this is to say that decompression isn’t our bottleneck. And this is despite decompression competing for GPU cores with the computation. The newer NVIDIA Blackwell Architecture includes a dedicated Decompression Engine which improves the decompression throughput even more.

And for those curious, a brief experiment without compression is about twice as slow on the GPU as the version with compression, though I didn’t investigate it deeply.

Computation

This example is primarily focused on the data loading portion of a Zarr workload, so the computation is secondary. I just threw in a bunch of matrix multiplications / reductions (which GPUs tend to do quickly).

But while the specific computation is unimportant, there are some characteristics to consider about your computation, it should take some non-negligible amount of time, such it’s worthwhile moving the data from the host to the device for the computation (and moving the result back to the host).

The key thing we care about here is overlapping host to device copies with compute, so that the GPU isn’t sitting around waiting for data. Note how the teal Host to Device Copy is running at the same time as the matrix multiplication from the previous shard:

And at this point, you can start analyzing GPU metrics if you still need to squeeze additional performance out of your pipeline.

But I think that’s enough for now.

Summary

One takeaway here is that GPUs are fast, which, sure. A slightly more interesting takeaway is that GPUs can be extremely fast, but achieving that takes some care.

In this workload my custom pipeline achieved high throughput by

1. Being very careful with memory allocations and data movement.
2. Using pinned host memory to speed up the one host to device transfer per shard
3. Use nvcomp and Zarr shards to parallelize decoding many chunks on the GPU
4. Use CUDA streams to express our workloads’ shard-level parallelism, so that we can overlap host I/O, host-to-device copies, kernel launches and kernel execution.

I’m hopeful that we can optimize the codec pipeline and memory handling in zarr-python to close the gap between what it provides and my custom, hand-optimized implementation (0.5s). But doing that in a general purpose library will require even more thought and care than my hacky implementation.

If you’ve made it this far, congrats. Reach out if you have any feedback, either directly or on the Zarr discussions board.

Last weekend I had the chance to sail in the 2025 Corn Coast Regatta. I had such a great time that I had to jot down my thoughts before they fade. This post is mostly for (future) me. We’ll return to our regularly scheduled programming in a future post. I have a post on Zarr performance cooking.

First, some context: in August I attended the Saylorville Yatch Club Sailing School Adult Small Boat class. This is a 3-day course that mixes some time in the classroom learning the theory and jargon (so much jargon!) with a bunch of time on the water. I had a bit of experience from sailing on summer weekends with my family growing up, but I wanted to learn more before going out on my own.

We were thrown in on the deep end, thanks to how breezy Saturday and Sunday. Too breezy for sailors as green as us, as it turns out. At least we got to practice capsize recovery a bunch.

Our instructor, Nick, was great. He’s knowledgeable, passionate about sailing, and invested in our success. If you’re near the area and at all interested in sailing, I’d recommend taking the course (and other clubs offer their own courses).

After the course, Nick was extremely generous. He invited us out for the Wednesday night beer can races the Yatch Club hosts, on his Melges 24. This was quite the step up from the Precision 185 we sailed during the class.

I hadn’t done any racing before and was immediately hooked. During the races, I was mostly just rail meat (“hiking” out on the lifelines to keep the boat from heeling over as we go upwind) and tried to not get in the way. But afterwards Nick was adamant about everyone getting to try the other jobs on the boat. Trimming the spinnaker (a very big sail that’s exclusively used going downwind) was awesome. There aren’t any winches on the Melges 24, so you directly feel the wind powering the boat when you’re flying the spinnaker.

Which brings us to last weekend. Nick was looking for some people to crew during the regatta, and we ended with Nick (driving), me (hiking upwind and flying the spin downwind), and a few other sailing school alumni on the boat. We started early on Saturday, rigging the special black carbon fiber sails Nick uses for regattas.

After a quick captains’ meeting we launched the boat and got ready to sail. Saturday was a series of short-distance buoy races. We ended up getting four races in, and our boat took second in each race to a Viper 640 sailed by a very experienced and talented father / son crew¹. A couple of races came down to the wire, and we might have won the third race if I hadn’t messed up our last gybe by grabbing the spinnaker sheet on the wrong side of the block and fouling everything up. Oops.

Sunday was the distance race. We started in about the middle of the lake, sailed a very wet 2 – 2.5 miles upwind (southeast) to the Saylorville Dam, followed by a long 4 – 5 mile downwind leg to the bridge on the north side of the lake, and finished with a ~2 mile leg to the end. The wind really picked up on Sunday, blowing ~15 kts with gusts up to 20–25 kts. As we neared the upwind mark, we had some discussions about whether or not to fly the spinnaker. That’s a lot of wind for a crew as inexperienced as we were (we’d only had one practice and the previous days’ races together). We took our time rounding the mark and eventually decided to set it. Nick took things easy on us, and overall things went well. We about went over twice (probably my fault; I was exhausted by the time we got 1/2 down the course) but our jib trimmer bailed us out both times just like we talked about in our pre-race talk. It sounds like even the Viper went over, so I don’t feel too bad.

Our team had really gelled by the end of the regatta. Crossing the finish line in first place was exhilarating. The official results aren’t posted yet, but we think we got first even after adjusting for the PHRF ratings.

I haven’t yet purchased (another) boat, but the Melges 15 and 19 both look fun (my poor old Honda CRV doesn’t have the towing capacity for a 24, alas). Regardless of what boat I’m on, I’m looking forward to spending more time on the water.

You can watch a video version of this talk at https://youtu.be/BFFHXNBj7nA

On Thursday, I presented a talk, GPU Accelerated Cloud-Native Geospatial, at the inaugural Cloud-Native Geospatial Conference (slides here). This post will give an overview of the talk and some background on the prep. But first I wanted to say a bit about the conference itself.

The organizers (Michelle Roby, Jed Sundell, and others from Radiant Earth) did a fantastic job putting on the event. I only have the smallest experience with helping run a conference, but I know it’s a ton of work. They did a great job hosting this first run of conference.

The conference was split into three tracks:

1. On-ramp to Cloud-Native Geospatial (organized by Dr. Julia Wagemann from thriveGEO)
2. Cloud-Native Geospatial in Practice (organized by Aimee Barciauskas from Development Seed)
3. Building Resilient Data Infrastructure Ecosystems (organized by Dr. Brianna Pagán, also from Development Seed)

Each of the track leaders did a great job programming their session. As tends to happen at these multi-track conferences, my only complaint is that there were too many interesting talks to choose from. Fortunately, the sessions were recorded and will be posted online. I spent most of my time bouncing between Cloud-Native Geospatial in Practice and On-ramp to Cloud-native Geospatial, but caught a couple talks from the Building Resilient Data Ecosystems track.

My main goal at the conference was to listen to peoples’ use-cases, with the hope of identifying workloads that might benefit from GPU optimization. If you have a geospatial workload that you want to GPU-optimize, please contact me.

My Talk

I pitched this talk about two months into my tenure at NVIDIA, which is to say about two months into my really using GPUs. In some ways, this made things awkward: here I am, by no means a CUDA expert, in front of a room telling people how they ought to be doing things. On the other hand, it’s a strength. I’m clearly not subject to the curse of expertise when it comes to GPUs, so I can empathize with what ended up being my intended audience: people who are new to GPUs and wondering if and where they can be useful for achieving their goals.

While preparing, I had some high hopes for doing deep-dives on a few geospatial workloads (e.g. Radiometric Terrain Correction for SAR data, pytorch / torchgeo / xbatcher dataloaders and preprocessing). But between the short talk duration, running out of prep time, and my general newness to GPUs, the talk ended up being fairly introductory and high-level. I think that’s OK.

GPUs are Fast

This was a fun little demo of a “quadratic means” example I took from the Pangeo forum. The hope was to get the room excited and impressed at just how fast GPUs can be. In it, we optimized the runtime of the computation from about 3 seconds on the CPU to about 20 ms on the GPU (via a one-line change to use CuPy).

For fun, we optimized it even further to just 4.5 ms by writing a hand-optimized CUDA to use some shared memory tricks and avoid repeated memory accesses.

You can see the full demo at https://github.com/TomAugspurger/gpu-cng-2025. I wish now that I had included more geospatial-focused demos. But the talk was only 15-20 minutes and already packed.

Getting Started with GPU programming

There is a ton of software written for NVIDIA chips. Before joining NVIDIA, I didn’t appreciate just how complex these chips are. NVIDIA, especially via RAPIDS, offers a bunch of relatively easy ways to get started.

This slide from Jacob Tomlinson’s PyData Global talk showcases the various “swim lanes” when it comes to programming NVIDIA chips from Python:

This built nicely off the demo, where we saw two of those swim lanes in action.

The other part lowering the barrier of entry is the cloud. Being programmable, a GPU is just an API call away (assuming you’re already set up on one of the clouds providing GPUs).

The I/O Problem

From there, we took a very high level overview of some geospatial workloads. Each loads some data (which we assumed came from Blob Storage), computed some result, and stored that result. For example, a cloud-free mosaic from some Sentinel-2 imagery:

I’m realizing now that I should have included a vector data example, perhaps loading an Overture Maps geoparquet file and doing a geospatial join.

Anyway, the point was to introduce some high-level concepts that we can use to identify workloads amenable to GPU acceleration. First, we looked at a workloads through time, which differ in how I/O vs. compute intensive they are.

For example, an I/O-bound workload:

Contrast that with a (mostly) CPU-bound workload:

Trying to GPU-accelerate the I/O-bound workload will only bring disappointment: even if you manage to speed up the compute portion, it’s such a small portion of the overall runtime to not make a meaningful difference.

But GPU-accelerating the compute-bound workload, on the other hand, can lead to to a nice speedup:

A few things are worth emphasizing:

1. You need to profile your workload to understand where time is being spent.
2. You might be able to turn an I/O bound problem into a compute bound problem by optimizing it (choosing a better file format, placing your compute next to the storage, choosing a faster library for I/O, parallelization, etc.)
3. I’m implying that “I/O” is just sitting around waiting on the network. In reality, some of I/O will be spent doing “compute” things (like parsing and decompressing bytes.) And those portions of I/O can be GPU accelerated.
4. I glossed over the “memory barrier” at this point in the talk, but returned to it later. There are again libraries (like KvikIO) that can help with this.

Pipelining

Some (most?) problems can be broken into smaller units of work and, potentially, parallelized. By breaking the larger problem into smaller pieces, we have the opportunity to optimize the throughput of our workload through pipelining.

Pipelining lets us overlap various parts of the workload that are using different parts of the system. For example I/O, which is mostly exercising the network, can be pipelined with computation, which is mostly exercising the GPU. First, we look at some poor pipelining:

The workload serially reads data, computes the result, and writes the output. This is inefficient: when you’re reading or writing data the GPU is idle (indeed, the CPU is mostly idle too, since it’s waiting for bytes to move over the network). And when you’re computing the result, the CPU (and network) are idle. This manifests as low utilization of the GPU, CPU, and network.

This second image shows good pipelining:

We’ve set up our program to read, compute, and write batches in parallel. We achieve high utilization of the GPU, CPU, and network.

This general concept can apply to CPU-only systems, especially multi-core systems. But the pain of low resource utilization is more pronounced with GPUs, which tend to be more expensive.

Now, this is a massively oversimplified example where the batches of work happen to be nicely sized and the workload doesn’t require an coordination across batches. But, with effort, the technique can be applied to a wide range of problems.

Memory Bandwidth

This section was pressed for time, but I really wanted to at least touch on one of the first things you’ll hit when doing data analysis on the GPU: moving data from host to device memory is relatively slow.

In the talk, I mostly just emphasized the benefits of leaving data on the GPU. The memory hierarchy diagram from the Flash Attention paper gave a nice visual representation of the tradeoff between bandwidth and size the different tiers give (I’d briefly mentioned the SRAM tier during the demo, since our most optimized version used SRAM).

But as I mentioned in the talk, most people won’t be interacting with the memory hierarchy beyond minimizing transfers between the host and device.

Reach Out

As I mentioned earlier my main goal attending the conference was to hear what the missing pieces of the GPU-accelerated geospatial landscape are (and to catch up with the wonderful members of this community). Reach out with any feedback you might have.

I have a new post up at the NVIDIA technical blog on High-Performance Remote IO with NVIDIA KvikIO.¹

This is mostly general-purpose advice on getting good performance out of cloud object stores (I guess I can’t get away from them), but has some specifics for people using NVIDIA GPUs.

In the RAPIDS context, NVIDIA KvikIO is notable because

1. It automatically chunks large requests into multiple smaller ones and makes those requests concurrently.
2. It can read efficiently into host or device memory, especially if GPU Direct Storage is enabled.
3. It’s fast.

As part of preparing this, I got to write some C++. Not a fan!

My local Department of Education has a public comment period for some proposed changes to Iowa’s science education standards. If you live in Iowa, I’d encourage you to read the proposal (PDF) and share feedback through the survey. If you, like me, get frustrated with how difficult it is to see what’s changed or link to a specific piece of text, read on.

I’d heard rumblings that there were some controversial changes around evolution and climate change. But rather than just believing what I read in a headline, I decided to do my own research (science in action, right?).

The proposed changes

I might have missed it, but I couldn’t find anywhere with the changes in an easily viewable form. The documents are available as PDFs (2015 standards, 2025 draft). The two PDFs aren’t formatted the same, making it very challenging to visually “diff” the two.

The programmers in the room will know that comparing two pieces of text is a pretty well solved problem. So I present to you, the changes:

The 2015 text is in red. The 2025 text is in green. That link includes just the top-level standards, not the “End of Grade Band Practice Clarification”, “Disciplinary Content Clarification”, or “End of Grade Band Conceptual Frame Clarification”.

The Python script I wrote to generate that diff took an hour or so to write and debug. If the standards had been in a format more accessible than a PDF it would have been minutes of work.

I’m somewhat sympathetic to the view that we should evaluate these new standards on their own terms, and not be biased by the previous language. But a quick glance at most of the changes shows you this is about language, and politics. It’s nice to be able to skim a single webpage to see that they’re just doing a Find and Replace for “evolution” and “climate change”.

Some thoughts

I’m mostly just disappointed. Disappointed in the people pushing this. Disappointed that they’re trying to claim the legitimacy of expertise

The standards were reviewed by a team consisting of elementary and secondary educators, administrators, content specialists, families, representatives from Iowa institutions of higher education and community partners.

and then saying they’re merely advisory

The team serves in an advisory capacity to the department; it does not finalize the first proposed revised draft standards

That’s a key component of pseudoscience: wrapping yourself in the language and of science and claiming expertise.

I’m disappointed that they they’re unwilling or unable to present the information in a easy to understand form.

I’m disappointed that they don’t live up to the documents’s own (well-put!) declaration on the importance of a good science education:

By the end of 12th grade, every student should appreciate the wonders of science and have the knowledge necessary to engage in meaningful discussions about scientific and technological issues that affect society. They must become discerning consumers of scientific and technological information and products in their daily lives.

Students’ science experiences should inspire and equip them for the reality that most careers require some level of scientific or technical expertise. This education is not just for those pursuing STEM fields; it’s essential for all students, regardless of their future education or career paths. Every Iowa student deserves an engaging, relevant, rigorous, and coherent pre-K–12 science education that prepares them for active citizenship, lifelong learning, and successful careers.

The survey includes a few questions about your overall feedback to the standards including, confusingly, a question asking if you agree or disagree that the standards will improve student learning, and then a required question asking you to “identify the reasons you believe that the recommended Iowa Science Standards will improve student learning”. I never took a survey design course, but it sure seems like I put more care into the pandas users surveys than this.

After answering the top-level questions about how great the new standards are, you have the option to provide specific feedback on each standard. Cheers to the people who actually go through each one and form an opinion. Mine focused on the ones that changed. I’ve included my responses below (header links go to the diff). Some extra commentary in the footnotes.

HS-LS2-7

A “Solution” is a solution to a problem. The proposed phrasing is awkward, and implies the need for “a solution to biodiversity”, i.e. that biodiversity is a problem that needs to be solved.

The previous text, “Design, evaluate, and refine a solution for reducing the impacts of human activities on the environment and biodiversity.” was clearer¹.

HS-LS4-1

The standard should make it clear that “biological change over time” refers specifically to “biological evolution”. Rephrase as

“Communicate scientific information that common ancestry and biological evolution are supported by multiple lines of empirical evidence.”²

HS-LS4-2

The standard should make clear that “biological change over time” is “evolution”. As Thomas Jefferson probably didn’t say, “The most valuable of all talents is that of never using two words when one will do.”³

HS-ESS2-3

I think there’s a typo somewhere in “cycling of matter magma”. Maybe “matter” was supposed to be replaced by “magma”?

HS-ESS2-4

The proposed standard seems to confuse stocks and flows, by saying that the flow of energy results into changes in climate trends. It’d be clearer to remove “trends”. If I dump 100 GJ of energy into a system, do I change its trend? No, unless you’re saying something about feedback effect and second derivatives (if so, make that clearer and focus on the feedback effects from global warming).

I recommend changing this to “Use a model to describe how variations in the flow of energy into and out of Earth’s systems result in changes in climate trends.”⁴

HS-ESS2-7

To make the interdependency between earth’s systems and life on earth clearer, I recommend phrasing this as “Construct an argument based on evidence about the simultaneous coevolution of Earth’s systems and life on Earth.”

This also gives our students a chance to learn the jargon they’ll hear, setting themselves up for success in the world.⁵

HS-ESS3-1

Phrasing this as “climate trends” narrows the standard to rule out abrupt changes in climate that aren’t necessarily part of a longer trend. I recommend phrasing this as “Construct an explanation based on evidence for how the availability of natural resources, occurrence of natural hazards, and changes in climate have influenced human civilizations.”

HS-ESS3-4

The proposed standard is unclear. It’s again using “solution” without stating what is being solved. What impact is being reduced?

Rephrase this as “Evaluate or refine a technological solution that reduces impacts of human activities on natural systems.”

HS-ESS3-5

Replace “climate trends” with “climate change”. We should ensure our students are ready for the language used in the field.

HS-ESS3-6

The standard should make it clear that human activity is the cause of the changes in the earth systems we’re currently experiencing. Rephrase the standard as “Use a computational representation to illustrate the relationships among Earth systems and how those relationships are being modified due to human activity.”

Again, if you’re in Iowa, read the proposals, check the diff, and leave feedback before February 3rd.

Over at https://github.com/opengeospatial/geoparquet/discussions/251, we’re having a nice discussion about how best to partition geoparquet files for serving over object storage. Thanks to geoparquet’s design, just being an extension of parquet, it immediately benefits from all the wisdom around how best to partition plain parquet datasets. The only additional wrinkle for geoparquet is, unsurprisingly, the geo component.

It’s pretty common for users to read all the features in a small spatial area (a city, say) so optimizing for that use case is a good default. Simplifying a bit, reading small spatial subsets of a larger dataset will be fastest if all the features that are geographically close together are also “close” together in the parquet dataset, and each part of the parquet dataset only contains data that’s physically close together. That gives you the data you want in the fewest number of file reads / HTTP requests, and minimizes the amount of “wasted” reads (data that’s read, only to be immediately discarded because it’s outside your area of interest).

Parquet datasets have two levels of nesting we can use to achieve our goal:

1. Parquet files within a dataset
2. Row groups within each parquet file

And (simplifying over some details again) we choose the number row groups and files so that stuff fits in memory when we actually read some data, while avoiding too many individual files to deal with.

So, given some table of geometries, we want to repartition (AKA shuffle) the records so that all the ones that are close in space are also close in the table. This process is called “spatial partitioning” or “spatial shuffling”.

Spatial Partitioning

Dewey Dunnington put together a nice post on various ways of doing this spatial partitioning on a real-world dataset using DuckDB. This post will show how something similar can be done with dask-geopandas.

Prep the data

A previous post from Dewy shows how to get the data. Once you’ve downloaded and unzipped the Flatgeobuf file, you can convert it to geoparquet with dask-geopandas.

The focus today is on repartitioning, not converting between file formats, so let’s just quickly convert that Flatgeobuf to geoparquet.

root = pathlib.Path("data")
info = pyogrio.read_info(root / "microsoft-buildings-point.fgb")
split = root / "microsoft-buildings-point-split.parquet"

n_features = info["features"]

CHUNK_SIZE=1_000_000
print(n_features // CHUNK_SIZE + 1)

chunks = dask.array.core.normalize_chunks((CHUNK_SIZE,), shape=(n_features,))
slices = [x[0] for x in dask.array.core.slices_from_chunks(chunks)]

def read_part(rows):
    return geopandas.read_file("data/microsoft-buildings-point.fgb", rows=rows)[["geometry"]]

df = dask.dataframe.from_map(read_part, slices)
shutil.rmtree(split, ignore_errors=True)
df.to_parquet(split, compression="zstd")

Spatial Partitioning with dask-geopandas

Now we can do the spatial partitioning with dask-geopandas. The dask-geopandas user guide includes a nice overview of the background and different options available. But the basic version is to use the spatial_shuffle method, which computes some good “divisions” of the data and rearranges the table to be sorted by those.

df = dask_geopandas.read_parquet(split)
%time shuffled = df.spatial_shuffle(by="hilbert")
%time shuffled.to_parquet("data/hilbert-16.parquet", compression="zstd")

On my local machine (iMac with a 8 CPU cores (16 hyper-threaded) and 40 GB of RAM), discovering the partitions took about 3min 40s. Rewriting the data to be shuffled took about 3min 25s. Recent versions of Dask include some nice stability and performance improvements, led by the folks at Coiled, which made this run without issue. I ran this locally, but it would be even faster (and scale to much larger datasets) with a cluster of machines and object-storage.

Now that they’re shuffled, we can plot the resulting spatial partitions:

r = dask_geopandas.read_parquet("data/hilbert-16.parquet")
ax = r.spatial_partitions.plot(edgecolor="black", cmap="tab20", alpha=0.25, figsize=(12, 9))
ax.set_axis_off()
ax.set(title="Hilbert partitioning (level=16)")

The outline of the United States is visible, and the spatial partitions do a good (but not perfect) job of making mostly non-overlapping, spatially compact partitions.

which gives

Here’s a similar plot for by="geohash"

And for by="morton"

Each partition ends up with approximately 1,000,000 rows (our original chunk size). Here’s a histogram of the count per partition:

import seaborn as sns
counts = [fragment.count_rows() for fragment in pyarrow.parquet.ParquetDataset("data/hilbert-16.parquet/").fragments]

sns.displot(counts);

The discussion also mentions KD trees as potentially better way of doing the partitioning. I’ll look into that and will follow up if anything comes out of it.

Here’s another Year in Books (I missed last year, but here’s 2022).

Most of these came from recommendations by friends, The Incomparable’s Book Club and (a new source), the “Books in the Box” episodes of Oxide and Friends.

The Soul of a New Machine, by Tracy Kidder

I technically read it in the last few days of 2023, but included here because I liked it so much. This came recommended by the Oxide and Friends podcast’s Books in the Box episode. I didn’t know a ton about the history of computing, but have been picking up an appreciation for it thanks to reading this book. It goes into a ton of detail about what it took Data General to design and release a new machine. Highly recommended to anyone interested in computing.

More Murderbot Diaries

I got caught up on Martha Well’s Murderbot Diaries series, finishing both Fugitive Telemetry and System Collapse. These continue to be so enjoyable. (This Wired piece about Martha Wells and the series is in my reading list).

Nona the Ninth, by Tamsyn Muir

This is third installment in her Locked Tomb series. I don’t remember a ton of details from the plot, but I do recall

1. This feeling very different from the previous entries (each of which felt different from their predecessors)
2. A general feeling of discomfort and tension, like things could explode at any time, which I think was deliberate

It’s not as simple to describe as “lesbian necromancers in space” like the Gideon the First, but overall, I enjoyed it.

The Cemeteries of Amalo series, by Katherine Addison

These are set in the same universe as The Goblin Emperor, but follow a different main character. I didn’t love these quite as much as The Goblin Emperor (which is just… perfect), but the writing in these is still great. Don’t expect a ton from the plot. These are still more about the world and characters moving through it than anything else.

The Hunt for Red October, by Tom Clancy

This is probably a sign that I’m entering middle age, but yeah this was a fun read. I think I picked this up after Bobby Chesney and Steve Vladek were reminiscing about Clancy novels on the NSL podcast. I didn’t make it through Patriot Games, though, so maybe I still have some youth in me?

Bookshops and Bonedust by Travis Baldree

This is a prequel to Legends & Lattes. If you enjoyed that, you’ll enjoy this one too.

Lord of the Rings by J.R.R. Tolkien

Me and my 8-year old have been working our way through these. We finished The Two Towers earlier in the year and will wrap up Return of the King this week. I’m not sure how much he appreciates all the detailed descriptions of the scenery, but he seems to be mostly following the plot. They continue to be perfect.

A Short History of Nearly Everything by Bill Bryson

I didn’t learn a ton of new actual science from this (humblebrag). If you have a decent high school or liberal arts education you’ll hopefully be familiar with most of the concepts. But I’d recommend reading it regardless because of all the background on the history and people involved in the discoveries (which my courses didn’t cover) and for the great writing. Also, I just love the idea of trying to cover everything in a single, general-audience book.

The Golden Enclaves by Naomi Novik

This is the third in the Scholomance trilogy. The first couple were great. The first especially was very fun, almost pop-corn fantasy (despite a lot of death. Like a lot). But this one somehow is way deeper, and in a way that makes you reevaluate the previous books. It’s maybe less “fun” because of where the story goes, but still great. I read this more recently but it’s stuck with me.

Jonathan Strange & Mr Norell by Susanna Clarke

This is a bit hard to review. It does seem to be long (I read it on a Kobo, but wow I see now that Goodreads says 1,006 pages). And while stuff happens, it’s not exactly action packed. Still, I never felt bored reading it, and I was able to follow things clearly the entire time. I think the characters were just so well written that she could bring back a character we haven’t heard from in 400 pages and have us immediately understand who they are and why they’re doing what they’re doing.

Susanna Clarke also wrote Piranesi which I still think about from time to time, and would highly recommend (despite even less happening in that book).

This is How You Lose the Time War by Amal El-Mohtar and Max Gladstone

This was a reread (I needed something short after the tome that was Jonathan Strange & Mr Norell), but this book had stuck with me since I first read it in 2021. It’s just so, so good. I guess it’s technically a romance set in a Sci-Fi world, which isn’t my usual genera. But I loved it mainly for the writing.

The setting is somewhat interesting, but that’s not really the point: two factions are in a struggle spanning multiple universes (“strands”, in the book). Their agents can travel through time and between strands, and embed themselves in various situations to nudge events along a favorable path. I love a good time-travel book, even if they don’t get into the mechanics.

The characters are somewhat interesting, but they’re also not really the point. We don’t get ton of detail about them (not even their real names; just get “Red” and “Blue”).

And the plot is also somewhat interesting, but I think still not the point. Stuff happens. They write letters to each other. More stuff happens. They fall in love. More stuff happens.

To me, it really comes down to the beautiful writing (with just enough structure around it to make all that flowery prose feel appropriate). I mean… just listen: “I distract myself. I talk of tactics and of methods. I say how I know how I know. I make metaphors to approach the enormous fact of you on slant.”

Overall, I’d recommend this to just about anyone. Plus, it’s short enough that it’s not a huge time commitment if it’s not your cup of tea.

Other

Some honorable, non-book mentions that I’ve started reading this year:

• The Money Stuff newsletter from Matt Levine for money stuff
• The One First newsletter from Steve Steve Vladeck for legal stuff
• Simon Willison’s blog for AI stuff
• Marc Brooker’s blog for distributed computing stuff

Overall, a solid year! My full list is on Goodreads Reach out to me if you have any questions or recommendations.

This post is a bit of a tutorial on serializing and deserializing Python dataclasses. I’ve been hacking on zarr-python-v3 a bit, which uses some dataclasses to represent some metadata objects. Those objects need to be serialized to and deserialized from JSON.

This is a (surprisingly?) challenging area, and there are several excellent libraries out there that you should probably use. My personal favorite is msgspec, but cattrs, pydantic, and pyserde are also options. But hopefully this can be helpful for understanding how those libraries work at a conceptual level (their exact implementations will look very different.) In zarr-python’s case, this didn’t quite warrant needing to bring in a dependency, so we rolled our own.

Like msgspec and cattrs, I like to have serialization logic separate from the core metadata logic. Ideally, you don’t need to pollute your object models with serialization methods, and don’t need to shoehorn your business logic to fit the needs of serialization (too much). And ideally the actual validation is done at the boundaries of your program, where you’re actually converting from the unstructured JSON to your structured models. Internal to your program, you have static type checking to ensure you’re passing around the appropriate types.

This is my first time diving into these topics, so if you spot anything that’s confusing or plain wrong, then let me know.

Overview

At a high level, we want a pair of methods that can serialize some dataclass instance into a format like JSON and deserialize that output back into the original dataclass.

The main challenge during serialization is encountering fields that Python’s json module doesn’t natively support. This might be “complex” objects like Python datetimes or NumPy dtype objects. Or it could be instances of other dataclasses if you have some nested data structure.

When deserializing, there are lots of pitfalls to avoid, but our main goal is to support typed deserialization. Any time we converted a value (like a datetime to a string, or a dataclass to a dict), we’ll need to undo that conversion into the proper type.

Example

To help make things clearer, we’ll work with this example:

@dataclasses.dataclass
class ArrayMetadata:
    shape: tuple[int, ...]
    timestamp: datetime.datetime  # note 1


@dataclasses.dataclass
class EncoderA:
    value: int

@dataclasses.dataclass
class EncoderB:
    value: int


@dataclasses.dataclass
class Metadata:
    version: typing.Literal["3"]   # note 2
    array_metadata: ArrayMetadata  # note 2
    encoder: EncoderA | EncoderB   # note 4
    attributes: dict[str, typing.Any]
    name: str | None = None     # note 5

Successfully serializing an instance of Metadata requires working through a few things:

1. Python datetimes are not natively serializable by Python’s JSON encoder.
2. version is a Literal["3"], in other words "3" is only valid value there. We’d ideally validate that when deserializing Metadata (since we can’t rely on a static linter like mypy to validate JSON data read from a file).
3. Metadata.array_metadata is a nested dataclass. We’ll need to recursively apply any special serialization / deserialization logic to any dataclasses we encounter
4. Metadata.encoder is a union type, between EncoderA and EncoderB. We’ll need to ensure that the serialized version has enough information to deserialize this into the correct variant of that Union
5. name is an Optional[str]. This is similar to a Union between two concrete types, where one of the types happens to be None.

Serialization

Serialization is relatively easy compared to deserialization. Given an instance of Metadata, we’ll use dataclasses.asdict to convert the dataclass to a dictionary of strings to values. The main challenge is telling the JSON encoder how to serialize each of those values, which might have be “complex” types (whether they be dataclasses or some builtin type like datetime.datetime). There are a few ways to do this, but the simplest way to do it is probably to use the default keyword of json.dumps.

def encode_value(x):
    if dataclasses.is_dataclass(x):
        return dataclasses.asdict(x)
    elif isinstance(x, datetime.datetime):
        return x.isoformat()
    # other special cases... 

    return x

If Python encounters a value it doesn’t know how to serialize, it will use your function.

>>> json.dumps({"a": datetime.datetime(2000, 1, 1)}, default=serialize)
'{"a": "2000-01-01T00:00:00"}'

For aesthetic reasons, we’ll use functools.singledispatch to write that:

import dataclasses, datetime, typing, json, functools


@functools.singledispatch
def encode_value(x: typing.Any) -> typing.Any:
    if dataclasses.is_dataclass(x):
        return dataclasses.asdict(x)

    return x

@encode_value.register(datetime.datetime)
@encode_value.register(datetime.date)
def _(x: datetime.date | datetime.datetime) -> str:
    return x.isoformat()


@encode_value.register(complex)
def _(x: complex) -> list[float, float]:
    return [x.real, x.imag]

# more implementations for additional type...

You’ll build up a list of supported types that your system can serialize.

And define your serializer like so:

def serialize(x):
    return json.dumps(x, default=encode_value)

and use it like:

>>> metadata = Metadata(
...     version="3",
...     array_metadata=ArrayMetadata(shape=(2, 2),
...     timestamp=datetime.datetime(2000, 1, 1)),
...     encoder=EncoderA(value=1),
...     attributes={"foo": "bar"}
... )
>>> serialized = serialize(metadata)
>>> serialized
'{"version": "3", "array_metadata": {"shape": [2, 2], "timestamp": "2000-01-01T00:00:00"}, "encoder": {"value": 1}, "attributes": {"foo": "bar"}, "name": null}'

Deserialization

We’ve done serialization, so we should be about halfway done, right? Ha! Because we’ve signed up for typed deserialization, which will let us faithfully round-trip some objects, we have more work to do.

A plain “roundtrip” like json.loads only gets us part of the way there:

>>> json.loads(serialized)
{'version': '3',
 'array_metadata': {'shape': [2, 2], 'timestamp': '2000-01-01T00:00:00'},
 'encoder': {'value': 1},
 'attributes': {'foo': 'bar'},
 'name': None}

We have plain dictionaries instead of instances of our dataclasses and the timestamp is still a string. In short, we need to decode all the values we encoded earlier. To do that, we need the user to give us a bit more information: We need to know the desired dataclass to deserialize into.

def deserialize(into: type[T], data: bytes) -> T:
    ...

Given some type T (which we’ll assume is a dataclass; we could do some things with type annotations to actually check that) like Metadata, we’ll build an instance using the deserialized data (with the properly decoded types!)

Users will call that like

>>> deserialize(into=Metadata, data=deserialized)
Metadata(...)

For a dataclass type like Metadata, we can get the types of all of its fields at runtime with typing.get_type_hints:

>>> typing.get_type_hints(Metadata)
{'version': typing.Literal['3'],
 'array_metadata': __main__.ArrayMetadata,
 'encoder': __main__.EncoderA | __main__.EncoderB,
 'attributes': dict[str, typing.Any],
 'name': str | None}

So we “just” need to write a decode_value function that mirrors our encode_value function from earlier.

def decode_value(into: type[T], value: Any) -> T:
    # the default implementation just calls the constructor, like int(x)
    # In practice, you have to deal with a lot more details like
    # Any, Literal, etc.
    return into(value)


@decode_value.register(datetime.datetime)
@decode_value.register(datetime.date)
def _(into, value):
    return into.fromisoformat(value)


@decode_value.register(complex)
def _(into, value):
    return into(*value)

# ... additional implementations

Unfortunately, “just” writing that decoder proved to be challenging (have I mentioned that you should be using msgspec for this yet?). Probably the biggest challenge was dealing with Union types. The msgspec docs cover this really well in its Tagged Unions section, but I’ll give a brief overview.

Let’s take a look at the declaration of encoder again:

@dataclasses.dataclass
class EncoderA:
    value: int

@dataclasses.dataclass
class EncoderB:
    key: str
    value: int


class Metadata:
    ...
    encoder: EncoderA | EncoderB

Right now, we serialize that as something like this:

{
    "encoder": {
        "value": 1
    }
}

With that, it’s impossible to choose between EncoderA and EncoderB without some heuristic like “pick the first one”, or “pick the first one that succeeds”. There’s just not enough information available to the decoder. The idea of a “tagged union” is to embed a bit more information in the serialized representation that lets the decoder know which to pick.

{
    "encoder": {
        "value": 1,
        "type": "EncoderA",
    }
}

Now when the decoder looks at the type hints it’ll see EncoderA | EncoderB as the options, and can pick EncoderA based on the type field in the serialized object. We have introduced a new complication, though: how do we get type in there in the first place?

There’s probably multiple ways, but I went with typing.Annotated. It’s not the most user-friendly, but it lets you put additional metadata on the type hints, which can be used for whatever you want. We’d require the user to specify the variants of the union types as something like

class Tag:
    ...

class EncoderA:
    value: int
    type: typing.Annotated[typing.Literal["a"], Tag] = "a"

class EncoderB:
    value: int
    key: str
    type: typing.Annotated[typing.Literal["b"], Tag] = "b"

(Other libraries might use something like the classes name as the value (by default) rather than requiring a single-valued Literal there.)

Now we have a type key that’ll show up in the serialized form. When our decoder encounters a union of types to deserialize into, it can inspect their types hints with include_extras:

>>> typing.get_type_hints(EncoderA, include_extras=True)
{'value': int,
 'type': typing.Annotated[typing.Literal['a'], <class '__main__.Tag'>]}

By walking each of those pairs, the decoder can figure out which value in type maps to which dataclass type:

>>> tags_to_types
{
    "a": EncoderA,
    "b": EncoderB,
}

Finally, given the object {"type": "a", "value": 1} it can pick the correct dataclass type to use. Then that can be fed through decode_value(EncoderA, value) to recursively decode all of its types properly.

Conclusion

There’s much more to doing this well that I’ve skipped over in the name of simplicity (validation, nested types like list[Metadata] or tuples, good error messages, performance, extensibility, …). Once again, you should probably be using msgspec for this. But at least now you might have a bit of an idea how these libraries work and how type annotations can be used at runtime in Python.

I wrote up a quick introduction to stac-geoparquet on the Cloud Native Geo blog with Kyle Barron and Chris Holmes.

The key takeaway:

STAC GeoParquet offers a very convenient and high-performance way to distribute large STAC collections, provided the items in that collection are pretty homogenous

Check out the project at http://github.com/stac-utils/stac-geoparquet.

I have, as they say, some personal news to share. On Monday I (along with some very talented teammates, see below if you’re hiring) was laid off from Microsoft as part of a reorganization. Like my Moving to Microsoft post, I wanted to jot down some of the things I got to work on.

For those of you wondering, the Planetary Computer project does continue, just without me.

Reflections

It should go without saying that all of this was a team effort. I’ve been incredibly fortunate to have great teammates over the years, but the team building out the Planetary Computer was especially fantastic. Just like before, this will be very self-centered and project-focused, overlooking all the other people and work that went into this.

I’m a bit uncomfortable with all the navel gazing, but I am glad I did the last one so here goes.

The Hub

Our initial vision for the Planetary Computer had four main components:

1. Data (the actual files in Blob Storage, ideally in cloud-optimized formats)
2. APIs (like the STAC API which make the data usable; using raster geospatial data without a STAC API feels barbaric now)
3. Compute
4. Applications (which package all the low level details into reports or tools that are useful to decision makers)

Initially, my primary responsibility on the team was to figure out “Compute”. Dan Morris had a nice line around “it shouldn’t require a PhD in remote sensing and a PhD in distributed computing to use this data.”

After fighting with Azure AD and RBAC roles for a few weeks, I had the initial version of the PC Hub up and running. This was a more-or-less stock version of the daskhub helm deployment with a few customizations.

Aside from occasionally updating the container images and banning crypto miners (stealing free compute to burn CPU cycles on a platform built for sustainability takes some hutzpah), that was mostly that. While the JupyterHub + Dask on Kubernetes model isn’t perfect for every use case, it solves a lot of problems. You might still have to know a bit about distributed computing in order to run a large computation, but at least our users didn’t have to fight with Kubernetes (just the Hub admin, me in this case).

Probably the most valuable aspect of the Hub was having a shared environment where anyone could easily run our Example Notebooks. We also ran several “cloud native geospatial” tutorials on one-off Hubs deployed for a conference.

This also gave the opportunity to sketch out an implementation of Yuvi’s kbatch proposal. I didn’t end up having time to follow up on the initial implementation, but I still think there’s room for a very simple way to submit batch Jobs to the same compute powering your interactive JupyterHub sessions.

stac-vrt

Very early on in project¹, we had an opportunity to present on the Planetary Computer to Kevin Scott and his team. Our presentation included a short demo applying a Land Use / Land Cover model to some NAIP data. While preparing that, I noticed that doing rioxarray.open_rasterio on a bunch of NAIP COGs was slow. Basically, GDAL had to make an HTTP request to read the COG metadata of each file.

After reading some GitHub issues and Pangeo discussions, I learned about using GDAL VRTs as a potential solution to the problem. Fortunately, our STAC items had all the information needed to build a VRT, and rioxarray already knew how to open VRTs. We just needed a tool to build that VRT. That was stac-vrt.

I say “was” because similar functionality is now (better) implemented in GDAL itself, stackstac, and odc-stac.

This taught me that STAC can be valuable beyond just searching for data. The metadata in the STAC items can be useful during analysis too. Also, as someone who grew up in the open-source Scientific Python Ecosystem, it felt neat to get tools like xarray and Dask in front of the CTO of Microsoft.

geoparquet

I had a very small hand in getting geoparquet started, connecting Chris Holmes with Joris van den Bossche and the geopandas / geoarrow group. Since then my contributions have been relatively minor, but at least for a while the Planetary Computer could claim to host the most geoparquet data (by count of datasets and volume) than anyone else. Overture Maps probably claims that title now, which is fantastic.

stac-geoparquet

Pretty early on, we had some users with demanding use-cases where the STAC API itself was becoming a bottleneck. We pulled some tricks to speed up their queries, but this showed us there was a need to provide bulk access to the STAC metadata, where the number of items in the result is very large.

With a quick afternoon hack, I got a prototype running that converted our STAC items (which live in a Postgres database) to geoparquet (technically, this predated geoparquet!). The generic pieces of that tooling are at https://github.com/stac-utils/stac-geoparquet/ now. Kyle Barron recently made some really nice improvements to the library (moving much of the actually processing down into Apache Arrow), and Pete Gadomski is working on a Rust implementation.

For the right workloads, serving large collections of STAC metadata through Parquet (or even better, Delta or Iceberg or some other table format) is indispensable.

Data Pipelines

These are less visible externally (except when they break), but a couple years ago I took on more responsibility for the data pipelines that keep data flowing into the Planetary Computer. Broadly speaking, this included

1. Getting data from upstream sources to Azure Blob Storage
2. Creating STAC Items for new data and ingesting them into the Postgres database

Building and maintaining these pipelines was… challenging. Our APIs or database would occasionally give us issues (especially under load). But the onboarding pipelines required a steady stream of attention, and would also blow up occasionally when the upstream data providers changed something. https://sre.google/sre-book/monitoring-distributed-systems/ is a really handy resource for thinking about how to monitor this type of system. This was a great chance to learn.

pc-id

Before we publicly launched the Planetary Computer, we didn’t have a good idea of how we would manage users. We knew that we wanted to role things out somewhat slowly (at least access to the Hub; the data and APIs might have always been anonymously available?). So we knew we needed some kind of sign-up systems, and some sort of identity system that could be used by both our API layer (built on Azure’s API Management service) and our Hub.

After throwing around some ideas (Azure AD B2C? Inviting beta users as Guests in the Microsoft Corp tenant?), I put together the sketch of a Django application that could be the Identity backend for both API Management and the Hub. Users would sign in with their Work or Personal Microsoft Accounts (in the Hub or API Management Dev Portal) and our ID application would check that the user was registered and approved.

We added a few bells and whistles to the Admin interface to speed up the approval process, and then more or less didn’t touch it aside from basic maintenance. Django is great. I am by no means a web developer, but it let us get started quickly on a solid foundation.

Other Highlights

• xstac: A library for creating STAC metadata for xarray datasets
• stac-table: A library for creating STAC metadata for geopandas GeoDataFrames
• Several stactools-packages, to create STAC metadata for specific datasets
• Many presentations on Cloud Native Geospatial, including at AMS and Cloud Native Geospatial Day

There’s lots of STAC here. I’d like to think that we had a hand in shaping how the STAC ecosystem works, especially for more “exotic” datasets like tables and data cubes in NetCDF or Zarr format.

What’s Next?

Last time around, I ended things with the exciting announcement that I was moving to Microsoft. This time… I don’t know! This is my first time not having a job lined up, so I’ll hope to spend some time finding the right thing to work on.

One thing I’m trying to figure out is how much to stock to place in the geospatial knowledge I’ve picked up over the last four years. I’ve spent a lot of time learning and thinking about geospatial things (though I still cant’t explain the difference between a CRS and Datum). There’s a lot of domain-specific knowledge needed to use these geospatial datasets (too much domain-specificity, in my opinion). We’ll see if that’s useful.

Like I mentioned above, I wasn’t the only one who was laid off. There are some really talented people on the job market, both more junior and more senior. If you’re looking for someone you can reach me at [email protected].

Thanks for reading!

Ned Batchelder recently shared Real-world match/case, showing a real example of Python’s Structural Pattern Matching. These real-world examples are a great complement to the tutorial, so I’ll share mine.

While working on some STAC + Kerchunk stuff, in this pull request I used the match statement to parse some nested objects:

for k, v in refs.items():
    match k.split("/"):
        case [".zgroup"]:
            # k = ".zgroup"
            item.properties["kerchunk:zgroup"] = json.loads(v)
        case [".zattrs"]:
            # k = ".zattrs"
            item.properties["kerchunk:zattrs"] = json.loads(v)
        case [variable, ".zarray"]:
            # k = "prcp/.zarray"
            if u := item.properties["cube:dimensions"].get(variable):
                u["kerchunk:zarray"] = json.loads(refs[k])
            elif u := item.properties["cube:variables"].get(variable):
                u["kerchunk:zarray"] = json.loads(refs[k])
        case [variable, ".zattrs"]:
            # k = "prcp/.zattrs"
            if u := item.properties["cube:dimensions"].get(variable):
                u["kerchunk:zattrs"] = json.loads(refs[k])
            elif u := item.properties["cube:variables"].get(variable):
                u["kerchunk:zattrs"] = json.loads(refs[k])
        case [variable, index]:
            # k = "prcp/0.0.0"
            if u := item.properties["cube:dimensions"].get(variable):
                u.setdefault("kerchunk:value", collections.defaultdict(dict))
                u["kerchunk:value"][index] = refs[k]
            elif u := item.properties["cube:variables"].get(variable):
                u.setdefault("kerchunk:value", collections.defaultdict(dict))
                u["kerchunk:value"][index] = refs[k]

The for loop is iterating over a set of Kerchunk references, which are essentially the keys for a Zarr group. The keys vary a bit. They could be:

1. Group metadata keys like .zgroup and .zattrs, which apply to the entire group.
2. Array metadata keys like prcp/.zarray or prcp/.zattrs (prcp is short for precipitation), which apply to an individual array in the group.
3. Chunk keys, like prcp/0.0.0, prcp/0.0.1, which indicate the chunk index in the n-dimensional array.

The whole point of this block of code is to update some other data (either the STAC item or the value referenced by the key). Between the different kinds of keys and the different actions we want to take for each kind of key, this seems like a pretty much ideal situation for structural pattern matching.

The subject of our match is k.split("/"):

match k.split("/")

Thanks to the Kerchunk specification, we know that any key should have exactly 0 or 1 /s in it, so we can define different cases to handle each.

Specific string literals have special meaning (like ".zgroup" and ".zarray") and control the key we want to update, so we handle all those first.

And the final case handles everything else: any data variable and index will match the

case [variable, index]

The ability to bind the values like variable = "prcp" and index = "0.0.0" makes updating the target data structure seamless.

Combine that with the walrus operator (the v:=), dict.setdefault, and collections.defaultdict, we get some pretty terse, clever code. Looking back at it a couple months later it’s probably bit too clever.

I wanted to share an update on a couple of developments in the STAC ecosystem that I’m excited about. It’s a great sign that even after 2 years after its initial release, the STAC ecosystem is still growing and improving how we can catalog, serve, and access geospatial data.

STAC and Geoparquet

A STAC API is a great way to query for data. But, like any API serving JSON, its throughput is limited. So in May 2022, the Planetary Computer team decided to export snapshots of our STAC database as geoparquet. Each STAC collection is exported as a Parquet dataset, where each record in the dataset is a STAC item. We pitched this as a way to do bulk queries over the data, where returning many and many pages of JSON would be slow (and expensive for our servers and database).

Looking at the commit history, the initial prototype was done over a couple of days. I wish I had my notes from our discussions, but this feels like the kind of thing that came out of an informal discussion like “This access pattern kind of sucks”, followed by “What if we …. ?”, and then “Let’s try it!¹”. And so we tried it, and it’s been great!

I think STAC as geoparquet can become a standard way to transfer STAC data in bulk. Chris Holmes has an open PR defining a specification for what the columns and types should be, which will help more tools than just that stac-geoparquet library interpret the data.

And Kyle Barron has an open PR making the stac-geoparquet library “arrow-native” by using Apache Arrow arrays and tables directly (via pyarrow), rather than pandas / geopandas. When I initially sketched out stac-geoparquet, it might have been just a bit early to do that. But given that we’re dealing with complicated, nested types (which isn’t NumPy’s strong suite) and we aren’t doing any analysis (which is pandas / NumPy’s strong suite), this will be a great way to move the data around.

Now I’m just hoping for a PostgreSQL ADBC adapter so that our PostGIS database can output the STAC items as Arrow memory. Then we can be all Arrow from the time the data leaves the database to the time we’re writing the parquet files.

STAC and Kerchunk

Kerchunk is, I think, going to see some widespread adoption over the next year or two. It’s a project (both a Python library and a specification) for putting a cloud-optimized veneer on top of non-cloud optimized data formats (like NetCDF / HDF5 and GRIB2).

Briefly, those file formats tend not to work great in the cloud because

1. In the cloud, we want to read files over the network (data are stored in Blob Storage, which is on a different machine than your compute). These file formats are pretty complicated, and can typically only be read by one library implemented in C / C++, which isn’t always able to read data over the network.
2. Reading the metadata (to build a structure like an xarray Dataset) tends to require reading many small pieces of data from many parts of the file. This is slow over the network, where each small read could translate to an HTTP request. On an SSD, seeking around the file to gather metadata is fine. Over the network, it’s slow.

Together, those mean that you aren’t able to easily load subsets of the data (even if the data are internally chunked!). You can’t load the metadata to do your filtering operations, and even if you could you might need to download the whole file just to throw away a bunch of data.

That’s where Kerchunk comes in. The idea is that the data provider can scan the files once ahead of time, extracting the Kerchunk indices, which include

1. The metadata (dimension names, coordinate values, attributes, etc.), letting you build a high-level object like an xarray.Dataset without needing any (additional) HTTP requests.
2. The byte offsets for each chunk of each data variable, letting you access arbitrary subsets of the data without needing to download and discard unnecessary data.

You store that metadata somewhere (in a JSON file, say) and users access the original NetCDF / GRIB2 data via that Kerchunk index file. You can even do metadata-only operations, like combining data variables from many files, or concatenating along a dimension to make a time series, without ever downloading the data.

We’ve had some experimental support for accessing a couple datasets hosted on the Planetary Computer via Kerchunk indices for a while now. We generated some indices and through them up in Blob Storage, including them as an asset in the STAC item. I’ve never really been happy with how how that works in practice, because of the extra hop from STAC to Kerchunk to the actual data.

I think that Kerchunk is just weird enough and hard enough to use that it can take time for users to feel comfortable with it. It’s hard to explain that if you want the data from this NetCDF file, you need to download this other JSON file, and then open that up with this other fsspec filesystem (no, not the Azure Blob Storage filesystem where the NetCDF and JSON files are, that’ll come later), and pass that result to the Zarr reader in xarray (no, the data isn’t stored in Zarr, we’re just using the Zarr API to access the data via the references…).

Those two additional levels of indirection (through a sidecar JSON file and then the Zarr reader via fsspec’s reference file system) are a real hurdle. So some of my teammate’s are working on storing the Kerchunk indices in the STAC items.

My goal is to enable an access pattern like this:

>>> import xarray as xr
>>> import pystac_client
>>> catalog = pystac_client.Client.open("https://planetarycomputer.microsoft.com/api/stac/v1")

>>> items = catalog.search(collections=["noaa-nwm"], datetime="2023-10-15", query=...)
>>> ds = xr.open_dataset(items, engine="stac")

Where the step from STAC to xarray / pandas / whatever is as easy with NetCDF or GRIB2 data as it is with COGs are Zarr data (thanks to projects like stackstac and odc-stac.) This is using ideas from Julia Signell’s xpystac library for that final layer, which would know how to translate the STAC items (with embedded Kerchunk references) into an xarray Dataset.

I just made an update to xstac, a library for creating STAC items for data that can be reresented as an xarray Datasets, to add support for embedding Kerchunk indices in a STAC item representing a dataset. The goal is to be “STAC-native” (by using things like the datacube extension), while still providing enough information for Kerchunk to do its thing. I’ll do a proper STAC extension later, but I want to get some real-world usage of it first.

I think this is similar in spirit to how Arraylake can store Kerchunk indices in their database, which hooks into their Zarr-compatible API.

The main concern here is that we’d blow up the size of the STAC items. That would bloat our database, slow down STAC queries and responses. But overall, I think it’s worth it for the ergonomics when it comes to loading the data. We’ll see.

Getting Involved

Reach out, either on GitHub or by email, if you’re interested in getting involved in any of these projects.

Last week, I was fortunate to attend Dave Beazley’s Rafting Trip course. The pretext of the course is to implement the Raft Consensus Algorithm.

I’ll post more about Raft, and the journey of implementing, it later. But in brief, Raft is an algorithm that lets a cluster of machines work together to reliably do something. If you had a service that needed to stay up (and stay consistent), even if some of the machines in the cluster went down, then you might want to use Raft.

Raft achieves this consensus and availability through log replication. A single node of the cluster is elected as the Leader, and all other nodes are Followers. The Leader interacts with clients to accept new commands (set x=41, or get y). The Leader notes these commands in its logs and sends them to the other nodes in the cluster. Once the logs have been replicated to a majority of the nodes in a cluster, the Leader can apply the command (actually doing it) and respond to the client. That’s the “normal operation” mode of Raft. Beyond that, much of the complexity of Raft comes from handling all the edge cases (what if a leader crashes? What if the leader comes back? What if there’s a network partition and two nodes try to become leader? and on, and on)

Raft was just about perfect for a week-long course. It’s a complex enough problem to challenge just about anyone. But it’s not so big that a person can’t hope to implement (much of) it in a week.

I liked the structure of the course itself. The actual “lecture” time was pretty short. We’d typically start the day with a short overview of one component of the problem. But after that, we spent a majority of the time actually working on the project. Dave didn’t just throw us to the wolves, but there was many a reference to “Draw the rest of the owl”.

That said, I really benefited from Dave’s gentle nudges on which part of the puzzle to work on next. The design space of implementing Raft is incredibly large. A typical Raft implementation will need to handle, at a minimum:

1. Communicating between multiple machines
2. Handling events (messages over the network, timers to call elections, etc.)
3. Leader elections
4. The Log
5. Log replication
6. Achieving consensus
7. The State Machine (e.g. updating the Key Value store)
8. Client interaction (a surprisingly tricky part, that completely blew up my implementation)
9. Persistence

You can implement these in just about any order. Going into the class I had no idea which would be “best” to do first (I still don’t think there’s a right order, but focusing on the Log and Log replication does seem like as good a start as any).

And that’s just the order you do things in. There’s also the question of how you go about implementing it. Are you using threads and queues, or asyncio? Mutable or immutable data structures? How do you test and monitor this?

But I think the biggest decision is around how you actually architect the system. How do you break this large problem down into smaller components? And how do those components interact? That’s the kind of thinking that’s helpful in my day job, and this project really taught me a lot (specifically, that I still have a ton to learn about designing and implementing this type of system). Also, it reinforced how difficult distributed systems can be.

Our class was in-person (Dave’s last course in this specific office). While I missed my big monitor and fancy ergonomic keyboard of my home-office, (not to mention my family), I am glad I got to go in person. It was nice to just let out an exasperated sigh and chat with classmate about how they’re handling a particularly tricky part of the project. The loved the informal conversations at breakfast and lunch (which inevitably turned back to aft).

I want to clean up a few parts of my implementation (AKA, trash the whole thing and start over). Once done I’ll make a followup post.

Thanks to Dave for hosting a great course, the other classmates, and to my family for letting me ditch them to go type on a laptop for a week.

A few colleagues and I recently presented at the CIROH Training and Developers Conference. In preparation for that I created a Jupyter Book. You can view it at https://tomaugspurger.net/noaa-nwm/intro.html I created a few cloud-optimized versions for subsets of the data, but those will be going away since we don’t have operational pipelines to keep them up to date. But hopefully the static notebooks are still helpful.

Lessons learned

Aside from running out of time (I always prepare too much material for the amount of time), I think things went well. JupyterHub (perhaps + Dask) and Kubernetes continues to be a great way to run a workshop.

The code for processing the data into cloud-optimized formats (either Kerchunk indexes, Zarr, or (geo)parquet) is at https://github.com/TomAugspurger/noaa-nwm/tree/main/processing

To process the data I needed to create some Dask clusters. I had the opportunity to use dask-kubernetes’ new Dask Operator. It was great!

The actual pipelines for processing the raw files into cloud-optimized formats (or Kerchunk indexes) continues to be a challenge. A large chunk of that complexity does come from the data itself, and I gather that the National Water Model is pretty complex, at a fundamental level. I ran into issues with corrupt files (which have since been fixed). An update to the National Water Model changed its internal chunking structure, which is incompatible with Kerchunk’s current implementation. These were pretty difficult to debug.

I think the main takeway from the conference was that we (either the users of this data, the Planetary Computer, NODD, or the Office of Water Prediction) needs to do something to make this data more usable on the cloud. Most likely some sort of Kerchunk index is the first stop, but this won’t handle every use case (see the timeseries notebook for an example). Maintaining operational pipelines is a challenge, but hoepfully we can take it on some day.

Over in Planetary Computer land, we’re working on bringing Sentinel-5P into our STAC catalog.

STAC items require a geometry property, a GeoJSON object that describes the footprint of the assets. Thanks to the satellites’ orbit and the (spatial) size of the assets, we started with some…interesting… footprints:

That initial footprint, shown in orange, would render the STAC collection essentially useless for spatial searches. The assets don’t actually cover (most of) the southern hemisphere.

Pete Gadomski did some really great work to understand the problem and fix it (hopefully once and for all). As the satellite crosses the antimeridian, a pole, or both, naive approaches to generating a footprint fails. It takes some more complicated logic to generate a good geometry. That’s now available as antimeridian on PyPI. It produces much more sensible footprints:

Building Tools

The real reason I wanted to write this post was to talk about tool building. This is a common theme of the Oxide and Friends podcast, but I think spending time building these kinds of small, focused tools almost always pays off.

Pete had a handful of pathologic test cases in the antimeridian test suite, but I wanted a way to quickly examine hundreds of footprints that I got back from our test STAC catalog. There are probably already tools for this, but I was able to put one together in Jupyter in about 10 minutes by building on Jupyter Widgets and ipyleaflet.

You can see it in action here (using Sentinel-2 footprints rather than Sentinel 5-P):

We get a STAC footprint browser (connected to our Python kernel!) with a single, pretty simple function.

m = ipyleaflet.Map(zoom=3)
m.layout.width = "600px"
layer = ipyleaflet.GeoJSON()
m.add(layer)


@ipywidgets.interact(item: pystac.ItemCollection = items)
def browse(item: pystac.Item):
    shape = shapely.geometry.shape(item)
    m.center = tuple(shape.centroid.coords[0])[::-1]

    layer.data = item.geometry
    print(item.id, item.datetime.isoformat())

Using this browser, I could quickly scrub through the Sentinel-5P items with the arrow keys and verify that the footprints looked reasonable.

The demo for this lives in the Planetary Computer Examples repository, and you can view the rendered version.

Today, I was debugging a hanging task in Azure Batch. This short post records how I used py-spy to investigate the problem.

Background

Azure Batch is a compute service that we use to run container workloads. In this case, we start up a container that processes a bunch of GOES-GLM data to create STAC items for the Planetary Computer . The workflow is essentially a big

for url in urls:
    local_file = download_url(url)
    stac.create_item(local_file)

We noticed that some Azure Batch tasks were hanging. Based on our logs, we knew it was somewhere in that for loop, but couldn’t determine exactly where things were hanging. The goes-glm stactools package we used does read a NetCDF file, and my experience with Dask biased me towards thinking the netcdf library (or the HDF5 reader it uses) was hanging. But I wanted to confirm that before trying to implement a fix.

Debugging

I wasn’t able to reproduce the hanging locally, so I needed some way to debug the actual hanging process itself. My go-to tool for this type of task is py-spy. It does a lot, but in this case we’ll use py-spy dump to get something like a traceback for what’s currently running (and hanging) in the process.

Azure Batch has a handy feature for SSH-ing into the running task nodes. With an auto-generated user and password, I had a shell on the node with the hanging process.

The only wrinkle here is that we’re using containerized workloads, so the actual process was in a Docker container and not in the host’s process list (I’ll try to follow Jacob Tomlinson’s lead and be intentional about container terminology). The py-spy documentation has some details on how to use py-spy with docker. This comment in particular has some more details on how to run py-spy on the host to detect a process running in a container. The upshot is a command like this, run on the Azure Batch node:

$ root@...:/home/yqjjaq/# docker run -it --pid=container:244fdfc65349 --rm --privileged --cap-add SYS_PTRACE python /bin/bash

where 244fdfc65349 is the ID of the container with the hanging process. I used the python image and then pip installed py-spy in that debugging container (you could also use some container image with py-spy already installed). Finally, I was able to run py-spy dump inside that running container to get the trace:

root@306ad36c7ae3:/# py-spy dump --pid 1
Process 1: /opt/conda/bin/python /opt/conda/bin/pctasks task run blob://pctaskscommon/taskio/run/827e3fa4-be68-49c9-b8c3-3d63b31962ba/process-chunk/3/create-items/input --sas-token ... --account-url https://pctaskscommon.blob.core.windows.net/
Python v3.8.16 (/opt/conda/bin/python3.8)

Thread 0x7F8C69A78740 (active): "MainThread"
    read (ssl.py:1099)
    recv_into (ssl.py:1241)
    readinto (socket.py:669)
    _read_status (http/client.py:277)
    begin (http/client.py:316)
    getresponse (http/client.py:1348)
    _make_request (urllib3/connectionpool.py:444)
    urlopen (urllib3/connectionpool.py:703)
    send (requests/adapters.py:489)
    send (requests/sessions.py:701)
    request (requests/sessions.py:587)
    send (core/pipeline/transport/_requests_basic.py:338)
    send (blob/_shared/base_client.py:333)
    send (blob/_shared/base_client.py:333)
    send (core/pipeline/_base.py:100)
    send (core/pipeline/_base.py:69)
    send (core/pipeline/_base.py:69)
    send (blob/_shared/policies.py:290)
    send (core/pipeline/_base.py:69)
    send (core/pipeline/_base.py:69)
    send (core/pipeline/_base.py:69)
    send (blob/_shared/policies.py:489)
    send (core/pipeline/_base.py:69)
    send (core/pipeline/policies/_redirect.py:160)
    send (core/pipeline/_base.py:69)
    send (core/pipeline/_base.py:69)
    send (core/pipeline/_base.py:69)
    send (core/pipeline/_base.py:69)
    send (core/pipeline/_base.py:69)
    run (core/pipeline/_base.py:205)
    download (blob/_generated/operations/_blob_operations.py:180)
    _initial_request (blob/_download.py:386)
    __init__ (blob/_download.py:349)
    download_blob (blob/_blob_client.py:848)
    wrapper_use_tracer (core/tracing/decorator.py:78)
    <lambda> (core/storage/blob.py:514)
    with_backoff (core/utils/backoff.py:136)
    download_file (core/storage/blob.py:513)
    create_item (goes_glm.py:32)
    create_items (dataset/items/task.py:117)
    run (dataset/items/task.py:153)
    parse_and_run (task/task.py:53)
    run_task (task/run.py:138)
    run_cmd (task/_cli.py:32)
    run_cmd (task/cli.py:50)
    new_func (click/decorators.py:26)
    invoke (click/core.py:760)
    invoke (click/core.py:1404)
    invoke (click/core.py:1657)
    invoke (click/core.py:1657)
    main (click/core.py:1055)
    __call__ (click/core.py:1130)
    cli (cli/cli.py:140)
    <module> (pctasks:8)
Thread 0x7F8C4A84F700 (idle): "fsspecIO"
    select (selectors.py:468)
    _run_once (asyncio/base_events.py:1823)
    run_forever (asyncio/base_events.py:570)
    run (threading.py:870)
    _bootstrap_inner (threading.py:932)
    _bootstrap (threading.py:890)
Thread 0x7F8C4A00E700 (active): "ThreadPoolExecutor-0_0"
    _worker (concurrent/futures/thread.py:78)
    run (threading.py:870)
    _bootstrap_inner (threading.py:932)
    _bootstrap (threading.py:890)

And we’ve found our culprit! The line

download_file (core/storage/blob.py:513)

and everything above it indicates that the process is hanging in the download stage, not the NetCDF reading stage!

This fix

“Fixing” this is pretty easy. The Python SDK for Azure Blob Storage includes the option to set a read_timeout when creating the connection client. Now if the download hangs it should raise a TimeoutError. Then our handler will automatically catch and retry it, and hopefully succeed. It doesn’t address the actual cause of something deep inside the networking stack hanging, but it’s good enough for our purposes.

Update: 2023-02-28. Turns out, the “fix” wasn’t actually a fix. The process hung again the next day. Naturally, I turned to this blog post to find the incantations to run (which is why I wrote it in the first place).

As for getting closer to an actual cause of the hang, a colleague suggested upgrading Python versions since there were some fixes in that area between 3.8 and 3.11. After about a week, there have been zero hangs on Python 3.11.

The Planetary Computer made its January 2023 release a couple weeks back.

The flagship new feature is a really cool new ability to visualize the Microsoft AI-detected Buildings Footprints dataset. Here’s a little demo made by my teammate, Rob:

Currently, enabling this feature required converting the data from its native geoparquet to a lot of protobuf files with Tippecanoe. I’m very excited about projects to visualize the geoparquet data directly (see Kyle Barron’s demo) but for now we needed to do the conversion.

Hats off to Matt McFarland, who did the work on the data conversion and the frontend to support the rendering.

New Datasets

As usual, we have a handful of new datasets hosted on the Planetary Computer. Follow the link on each of these to find out more.

Climate Change Initiative Land Cover

NOAA Climate Normals¹

USDA Cropland Data Layer

USGS Land Change Monitoring, Assessment, and Projection

National Wetlands Inventory

Other stuff

We’ve also been doing a lot of work around the edges that doesn’t show up in visual things like new features or datasets. That work should show up in the next release and I’ll be blogging more about it then.

Over on the Planetary Computer team, we get to have a lot of fun discussions about doing geospatial data analysis on the cloud. This post summarizes some work we did, and the (I think) interesting conversations that came out of it.

Background: GOES-GLM

The instigator in this case was onboarding a new dataset to the Planetary Computer, GOES-GLM. GOES is a set of geostationary weather satellites operated by NOAA, and GLM is the Geostationary Lightning Mapper, an instrument on the satellites that’s used to monitor lightning. It produces some really neat (and valuable) data.

The data makes its way to Azure via the NOAA Open Data Dissemination program (NODD) as a bunch of NetCDF files. Lightning is fast [citation needed], so the GOES-GLM team does some clever things to build up a hierarchy of “events”, “groups”, and “flashes” that can all be grouped in a file. This happens very quickly after the data is captured, and it’s delivered to Azure soon after that. All the details are at https://www.star.nesdis.noaa.gov/goesr/documents/ATBDs/Baseline/ATBD_GOES-R_GLM_v3.0_Jul2012.pdf for the curious.

Cloud-native NetCDF?

The raw data are delivered as a bunch of NetCDF4 files, which famously isn’t cloud-native. The metadata tends to be spread out across the file, requiring many (small) reads to load the metadata. If you only care about a small subset of the data, those metadata reads can dominate your processing time. Remember: reading a new chunk of metadata typically requires another HTTP call. Even when your compute is in the same region as the data, an HTTP call is much slower than seeking to a new spot in an open file on disk.

But what if I told you that you could read all the metadata in a single HTTP request? Well, that’s possible with these NetCDF files. Not because of anything special about how the metadata is written, just that these files are relatively small. They’re only about 100-300 KB in total. So we can read all the metadata (and data) in a single HTTP call.

That gets to a point made by Paul Ramsey in his Cloud Optimized Shape File article:

One of the quiet secrets of the “cloud optimized” geospatial world is that, while all the attention is placed on the formats, the actual really really hard part is writing the clients that can efficiently make use of the carefully organized bytes.

So yes, the file formats do (often) matter. And yes, we need clients that can make efficient use of those carefully organized bytes. But when the files are this small, it doesn’t really matter how the bytes or organized. You’re still making a single HTTP call, whether you want all the data or just some of it.

This was a fun conversation amongst the team. We like to say we host “cloud-optimized data” on the Planetary Computer, and we do. But what really matters is the user experience. It’s all about the cloud-optimized vibes.

Build with users, not just with users in mind

A last, small point is the importance of getting user feedback before you go off doing something. We looked at the data and noticed the obviously tabular nature of the data and decided to split these single NetCDF file into three geoparquet files. In the abstract this make sense: these are naturally tabular, and parquet is the natural file format for them. We figured our users would appreciate the conversion. However we suddenly tripled the number of objects in Blob Storage. With this many objects and with new objects arriving so frequently, the sheer number of small files became a challenge to work with. This is, I think, still the right format for the data. But we’ll need to do more with our users to confirm that that’s the case before committing to maintain this challenging data pipeline to do the conversion at scale.

I came across a couple of new (to me) uses of queues recently. When I came up with the title to this article I knew I had to write them up together.

Queues in Dask

Over at the Coiled Blog, Gabe Joseph has a nice post summarizing a huge amount of effort addressing a problem that’s been vexing demanding Dask users for years. The main symptom of the problem was unexpectedly high memory usage on workers, leading to crashing workers (which in turn caused even more network communication, and so more memory usage, and more crashing workers). This is actually a problem I worked on a bit back in 2019, and I made very little progress.

A common source of this problem was having many (mostly) independent “chains” of computation. Dask would start on too many of the “root” tasks simultaneously, before finishing up some of the chains. The root tasks are typically memory increasing (e.g. load data from file system) while the later tasks are typically memory decreasing (take the mean of a large array).

In dask/distributed, Dask actually has two places where it determines which order to run things in. First, there’s a “static” ordering (implemented in dask/order.py, which has some truly great docstrings, check out the source.) Dask was actually doing really well here. Consider this task graph from the issue:

The “root” tasks are on the left (marked 0, 3, 11, 14). Dask’s typical depth-first algorithm works well here: we execute the first two root tasks (0 and 3) to finish up the first “chain” of computation (the box (0, 0) on the right) before moving onto the other two root nodes, 11 and 14.

The second time Dask (specifically, the distributed scheduler) considers what order to run things is at runtime. It gets this “static” ordering from dask.order which says what order you should run things in, but the distributed runtime has way more information available to it that it can use to influence its scheduling decisions. In this case, the distributed scheduler looked around and saw that it had some idle cores. It thought “hey, I have a bunch of these root tasks ready to run”, and scheduled those. Those tend to increase memory usage, leading to our memory problems.

The solution was a queue. From Gabe’s blog post:

We’re calling this mode of scheduling “queuing”, or “root task withholding”. The scheduler puts data-loading tasks in an internal queue, and only drips one out to a worker once it’s finished its current work and there’s nothing more useful to run instead that utilizes the work it just completed.

Queue for Data Pipelines

At work, we’re taking on more responsibility for the data pipeline responsible for getting various datasets to Azure Blob Storage. I’m dipping my toes into the whole “event-driven” architecture thing, and have become paranoid about dropping work. The Azure Architecture Center has a bunch of useful articles here. This article gives some names to some of the concepts I was bumbling through (e.g. “at least once processing”).

In our case, we’re using Azure Queue Storage as a simple way to reliably parallelize work across some machines. We somehow discover some assets to be copied (perhaps by querying an API on a schedule, or by listening to some events on a webhook), store those as messages on the queue.

Then our workers can start processing those messages from the queue in parallel. The really neat thing about Azure’s Storage Queues (and, I gather, many queue systems) is the concept of “locking” a message. When the worker is ready, it receives a message from the queue and begins processing it. To prevent dropping messages (if, e.g. the worker dies mid-processing) the message isn’t actually deleted until the worker tells the queue service “OK, I’m doing processing this message”. If for whatever reason the worker doesn’t phone home saying it’s processed the message, the message reappears on the queue for some other worker to process.

The Azure SDK for Python actually does a really good job integrating language features into the clients for these services. In this case, we can just treat the Queue service as an iterator.

>>> queue_client = azure.storage.blob.QueueClient("https://queue-endpoint.queue.core.windows.net/queue-name")
>>> for message in queue_client.receive_messages():
...    yield message
...    # The caller finishes processing the message.
...    queue_client.delete_message(message)

I briefly went down a dead-end solution that added a “processing” state to our state database. Workers were responsible for updating the items state to “processing” as soon as they started, and “copied” or “failed” when they finished. But I quickly ran into issues where items were marked as “processing” but weren’t actually. Maybe the node was preempted; maybe (just maybe) there was a bug in my code. But for whatever reason I couldn’t trust the item’s state anymore. Queues were an elegant way to ensure that we processed these messages at least once, and now I can sleep comfortably at night knowing that we aren’t dropping messages on the floor.

It’s “Year in X” time, and here’s my 2022 Year in Books on GoodReads. I’ll cover some highlights here.

Many of these recommendations came from the Incomparable’s Book Club, part of the main Incomparable podcast. In particular, episode 600 The Machine was a Vampire which is a roundup of their favorites from the 2010s.

Bookended by Murderbot Diaries

I started and ended this year (so far) with a couple installments in the Murderbot Diaries. These follow a robotic / organic “Security Unit” that’s responsible for taking care of humans in dangerous situations. We pick up after an unfortunate incident where it seems to have gone rouge and murdered her clients (hence, the murderbot) and hacked its governor module to essentially become “free”.

There’s some exploration of “what does it mean to be human?” in these, but mostly they’re just fun.

Competency

I read a pair of books this year that are set completely different worlds (one in some facsimile of the Byzantine empire, and another in the earth’s near-future) that are related by the protagonist being competent at engineering and problem solving.

First up was Andy Weir’s Project Hail Mary (a followup to The Martin, which falls under this category too). At times it felt like some challenges were thrown up just so that the main character could knock them down. But it also had one of my favorite fictional characters ever (no spoilers, but it’s Rocky).

The second was K.J. Parker’s Sixteen Ways to Defend a Walled City. In this one, the main character feels a bit more balanced. His strengths around engineering and problem solving are offset by his (self-admitted) weaknesses. I really enjoyed this one.

Some Classics

After reading Jo Walton’s Among Others, which follows a Sci-Fi / Fantasy obsessed girl as she goes through some… things, I dipped in to some of the referenced works I had never gotten to before.

First was Ursula K. Le Guin’s The Left Hand of Darkness. This was great. I imagine it was groundbreaking and controversial when it first came out, but I still liked it as a story.

Next was Kurt Vonnegut’s Cat’s Cradle. Wow, was this good. I’d only read Slaughterhouse-Five before, and finally got around to some of his other stuff. Sooo good.

Wholesomeness

There were two books that I just loved (both got 5 stars on goodreads) that I want to label “wholesome”.

Piranesi, by Susanna Clarke, was just great. The setup is bizarre, but we follow our… wholly innocent (naive? definitely wholesome) main character in a world of classical Greek statues and water. Piranesi just Loves his World and that’s great.

Next up is Katherine Addison’s The Goblin Emperor. This a story of a fundamentally good person unexpectedly thrown into power. He does not simply roll over and get pushed around by the system, and he retains his fundamental goodness. It’s pretty long (449 pages) and not much actually “happens” (there’s maybe two or three “action” scenes). And yet somehow Katherine kept the story moving and all the factions straight.

Not-so-wholesome

My other 5-star book this year was Cormac McCarthy’s The Road. I know it’s super popular so you don’t need me recommending it, but dang this got to me a bit¹. I don’t know how old The Boy is in the story, but mine’s six now and it was hard not to let imagination wander.

Nonfiction

I think the only non-fiction books I read this year were

• Command and Control by Eric Schlosser about how (not) to safely have a nuclear weapons arsenal
• The Fifth Risk: Undoing Democracy by Michael Lewis, about the dangers posed by putting people in government who don’t care about doing a good job
• The Path Between the Seas about the development and construction of the Panama canal

This is less than I would have liked, but hey, I’ve been tired.

The Rest

You can find my read books on goodreads. I don’t think I read (or at least finished) any bad books this year. My lowest-rated was Eye of the World (the first book in the Wheel of Time series) and it was… long. It world seems neat though. Leviathan Falls wrapped up the Expanse series satisfyingly. The Nova Incident is a fun spy / cold-war thriller set in the far future, which I’d recommend reading after the earlier ones in that series. On the other hand, Galaxy and the Ground Within (book 4 in the Wayfarers series) worked just fine without having read the others.

Overall, a good year in books!

Mike Duncan is wrapping up his excellent Revolutions podcast. If you’re at all interested in history then now is a great time to pick it up. He takes the concept of “a revolution” and looks at it through the lens of a bunch of revolutions throughout history. The appendix episodes from the last few weeks have really tied things together, looking at whats common (and not) across all the revolutions covered in the series.

It’s hard to believe that this podcast started in 2013. I came over from Mike’s The History of Rome podcast (which started in 2007(!) I’m not sure when I got on that train, but it was in the manually sync podcasts to an iPod days). Congrats to Mike for a podcast well done!

Like some others, I’m getting back into blogging.

I’ll be “straying from my lane” and won’t just be writing about Python data libraries (though there will still be some of that). If you too would like to blog more, I’d encourge you to read Simon Willison’s What to blog About and Matt Rocklin’s Write Short Blogposts.

Because I’m me, I couldn’t just make a new post. I also had to switch static site generators, just becauase. All the old links, including my RSS feed, should continue to work. If you spot any issues, let me know (I think I’ve fixed at least one bug in the RSS feed, apologies for any spurious updates. But just in case, you might want to update your RSS links to http://tomaugspurger.net/index.xml).

Speaking of RSS, it’s not dead! I’ve been pleasently surprised to see new activity in feeds I’ve subscribed to for years. (If you’re curious, I use NetNewsWire for my reader).

Some personal news: Last Friday was my last day at Anaconda. Next week, I’m joining Microsoft’s AI for Earth team. This is a very bittersweet transition. While I loved working at Anaconda and all the great people there, I’m extremely excited about what I’ll be working on at Microsoft.

Reflections

I was inspired to write this section by Jim Crist’s post on a similar topic: https://jcristharif.com/farewell-to-anaconda.html. I’ll highlight some of the projects I worked on while at Anaconda. If you want to skip the navel gazing, skip down to what’s next.

1. This is self-serving and biased to over-emphasize my own role in each of these. None of these could be done without the other individuals on those teams, or the support of my family.
2. More companies should support open-source like Anaconda does: offer positions to the maintainers of open-source projects and see what they can do. Anaconda recently announced a program that makes it easier for more companies to support open-source.

pandas

If I had a primary responsibility at Anaconda, it was stewarding the pandas project. When I joined Anaconda in 2017, pandas was around the 0.20 release, and didn’t have much in the way of paid maintenace. By joining Anaconda I was fulfilling a dream: getting paid to work on open-source software. During my time at Anaconda, I was the pandas release manager for a handful of pandas releases, including pandas 1.0.

I think the most important code to come out of my work on pandas is the extension array interface. My post on the Anaconda Blog tells the full story, but this is a great example of a for-profit company (Anaconda) bringing together a funding source and an open-source project to accomplish something great for the community. As an existing member of the pandas community, I was able to leverage some trust that I’d built up over the years to propose a major change to the library. And thanks to Anaconda, we had the funding to realisitically pull (some of) it off. The work is still ongoing, but we’re gradually solving some of pandas’ longest-standing pain points (like the lack of an integer dtype with missing values).

But even more important that the code is probably pandas winning its first-ever funding through the CZI EOSS program. Thanks to Anaconda, I was able to dedicate the time to writing the proposal. This work funded

1. Maintenance, including Simon Hawkins managing the last few releases.
2. A native string dtype, based on Apache Arrow, for faster and more memory-efficient strings (coming in the next release or two)
3. Many improvements to the extension array interface

Now that I’m leaving Anaconda, I suspect my day-to-day involvement in pandas will drop off a bit. But I’ll still be around, hopefully focusing most on helping others work on pandas.

Oh, side-note, I’m extremely excited about the duplicate label handling coming to pandas 1.2.0. That was fun to work on and I think will solve some common pandas papercuts.

Dask

I started using Dask before I joined Anaconda. It exactly solved my needs at the time (I was working with datasets that were somewhat larger than the memory of the machine I had access to). I was thrilled to have more time for working on it along with others from Anaconda; I learned a ton from them.

My personal work mainly focused on ensuring that dask.dataframe continued to work well with (and benefit from) the most recent changes to pandas. I also kicked off the dask-ml project, which initially just started as a bit of documentation on the various projects in the “dask / machine learning” space (like distributed.joblib, dask-searchcv, dask-xgboost). Eventually this grew into a project of its own, which I’m reasonably happy with, even if most people don’t need distributed machine learning.

pymapd

pymapd is a Python library that implements the DB API spec for OmniSci (FKA MapD). For the most part, this project involved copying the choices made by sqlite3 or psycopg2 and applying them to. The really fun part of this project was working with Wes McKinney, Siu Kwan Lam, and others on the GPU and shared memory integration. Being able to query a database and get back zero-copy results as a DataFrame (possibly a GPU DataFrame using cuDF) really is neat.

ucx-py

ucx-py is a Python library for UCX, a high-performance networking library. This came out of work with NVIDIA and Dask, seeing how we could speed up performance on communication-bound workloads (UCX supports high-performance interfaces between devices like NVLink). Working on ucx-py was my first real foray into asyncio and networking. Fortunately, while this was a great learning experience for me, I suspect that very little of my code remains. Hopefully the early prototypes were able to hit some of the roadblocks the later attempts would have stumbled over. See this post for an overview of what that team has been able to accomplish recently.

Pangeo

Some time last year, after Matt Rocklin left for NVIDIA, I filled his spot on a NASA ACCESS grant funding work on the Pangeo project. Pangeo is a really interesting community. They’re a bunch of geoscientists trying to analyze large datasets using tools like xarray, Zarr, Dask, holoviz, and Jupyter. Naturally, they find rough edges in that workflow, and work to fix them. That might mean working with organizations like NASA to provide data in analysis-ready form. It might mean fixing bugs or performance issues in libraries like Dask. Being able to dedicate large chunks of time is crucial to solving these types of thorny problems, which often span many layers (e.g. using xarray to read data Zarr data from Google cloud storage involves something like eight Python libraries). While there’s still work to be done, this type of workflow is smoother than it was a couple years ago.

In addition to work on Dask itself, I was able to help out Pangeo in a few other ways:

1. I helped maintain pangeo’s JupyterHub deployments at pangeo-cloud-federation. (FYI, 2i2c is a new organization that’s purpose-built to do this kind of work).
2. I put together the daskhub Helm Chart, which Pangeo previously developed and maintained. It combines Dask Gateway’s and JupyterHub’s helm charts, along with experience from pangeo’s deployments, to deploy a multi-user JupyterHub deployment with scalable computation provided by Dask.
3. I helped with rechunker, a library that very specifically solves a problem that had vexxed pangeo’s community members for years.

Overall, working with the Pangeo folks has been incredibly rewarding. They’re taking the tools we know and love, and putting them together to build an extremely powerful, open architechture toolchain. I’ve been extremely lucky to work on this project. Which brings me to…

What's Next

As I mentioned up top, I’m joining the AI for Earth team at Microsoft. I’ll be helping them build tools and environments for distributed geospatial data processing! I’m really excited about this work. Working with the Pangeo community has been incredibly rewarding. I’m lookingo forward to doing even more of that.

P.S. we’re hiring!

As pandas’ documentation claims: pandas provides high-performance data structures. But how do we verify that the claim is correct? And how do we ensure that it stays correct over many releases. This post describes

1. pandas’ current setup for monitoring performance
2. My personal debugging strategy for understanding and fixing performance regressions when they occur.

I hope that the first section topic is useful for library maintainers and the second topic is generally useful for people writing performance-sensitive code.

Know thyself

The first rule of optimization is to measure first. It’s a common trap to think you know the performance of some code just from looking at it. The difficulty is compounded when you’re reviewing a diff in a pull request and you lack some important context. We use benchmarks to measure the performance of code.

There’s a strong analogy between using unit tests to verify the correctness of code and using benchmarks to verify its performance. Each gives us some confidence that an implementation behaves as expected and that refactors are not introducing regressions (in correctness or performance). And just as you use can use a test runner like unittest or pytest to organize and run unit tests, you can use a tool to organize and run benchmarks.

For that, pandas uses asv.

airspeed velocity (asv) is a tool for benchmarking Python packages over their lifetime. Runtime, memory consumption and even custom-computed values may be tracked. The results are displayed in an interactive web frontend that requires only a basic static webserver to host.

asv provides a structured way to write benchmarks. For example, pandas Series.isin benchmark looks roughly like

class IsIn:

    def setup(self):
        self.s = Series(np.random.randint(1, 10, 100000))
        self.values = [1, 2]

    def time_isin(self):
        self.s.isin(self.values)

There’s some setup, and then the benchmark method starting with time_. Using the asv CLI, benchmarks can be run for a specific commit with asv run <commit HASH>, or multiple commits can be compared with asv continuous <GIT RANGE>. Finally, asv will collect performance over time and can visualize the output. You can see pandas’ at https://pandas.pydata.org/speed/pandas/.

Detecting Regressions

asv is designed to be run continuously over a project’s lifetime. In theory, a pull request could be accompanied with an asv report demonstrating that the changes don’t introduce a performance regression. There are a few issues preventing pandas from doing that reliably however, which I’ll go into later.

Handling Regressions

Here’s a high-level overview of my debugging process when a performance regression is discovered (either by ASV detecting one or a user reporting a regression).

To make things concrete, we’ll walk through this recent pandas issue, where a slowdown was reported. User reports are often along the lines of

DataFrame.memory_usage is 100x slower in pandas 1.0 compared to 0.25

In this case, DataFrame.memory_usage was slower with object-dtypes and deep=True.

v1.0.3: memory_usage(deep=True) took 26.4566secs

v0.24.0: memory_usage(deep=True) took 6.0479secs

v0.23.4: memory_usage(deep=True) took 0.4633secs

The first thing to verify is that it’s purely a performance regression, and not a behavior change or bugfix, by ensuring that the outputs match between versions. Sometimes correctness requires sacrificing speed. In this example, we confirmed that the outputs from 0.24 and 1.0.3 matched, so we focused there.

Now that we have what seems like a legitimate slowdown, I’ll reproduce it locally. I’ll first activate environments for both the old and new versions (I use conda for this, one environment per version of pandas, but venv works as well assuming the error isn’t specific to a version of Python). Then I ensure that I can reproduce the slowdown.

In [1]: import pandas as pd

In [2]: df = pd.DataFrame({"A": list(range(10000))}, dtype=object)

In [3]: %timeit df.memory_usage(deep=True)
5.37 ms ± 201 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [4]: pd.__version__
Out[4]: '0.25.1'

versus

In [1]: import pandas as pd

In [2]: df = pd.DataFrame({"A": list(range(10000))}, dtype=object)

In [3]: %timeit df.memory_usage(deep=True)
17.5 ms ± 98.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [4]: pd.__version__
Out[4]: '1.0.1'

So we do have a slowdown, from 5.37ms -> 17.5ms on this example.

Once I’ve verified that the outputs match and the slowdown is real, I turn to snakeviz (created by Matt Davis, which measures performance at the function-level. For large enough slowdowns, the issue will jump out immediately with snakeviz.

From the snakeviz docs, these charts show

the fraction of time spent in a function is represented by the extent of a visualization element, either the width of a rectangle or the angular extent of an arc.

I prefer the “sunburst” / angular extent style, but either works.

In this case, I noticed that ~95% of the time was being spent in pandas._libs.lib.memory_usage_of_object, and most of that time was spent in PandasArray.__getitem__ in pandas 1.0.3. This is where a bit of pandas-specific knowledge comes in, but suffice to say, it looks fishy¹.

As an aside, to create and share these snakeviz profiles, I ran the output of the %snakeviz command through svstatic and uploaded that as a gist (using gist). I then pasted the “raw” URL to https://rawgit.org/ to get the URL embedded here as an iframe.

Line Profiling

With snakeviz, we’ve identified a function or two that’s slowing things down. If I need more details on why that’s function is slow, I’ll use line-profiler. In our example, we’ve identified a couple of functions, IndexOpsMixin.memory_usage and PandasArray.__getitem__ that could be inspected in detail.

You point line-profiler at one or more functions with -f and provide a statement to execute. It will measure things about each line in the function, including the number of times it’s hit and how long is spent on that line (per hit and total)

In  [9]: %load_ext line_profiler
In [10]: %lprun -f pd.core.base.IndexOpsMixin.memory_usage df.memory_usage(deep=True)
Total time: 0.034319 s
File: /Users/taugspurger/miniconda3/envs/pandas=1.0.1/lib/python3.8/site-packages/pandas/core/base.py
Function: memory_usage at line 1340

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
  1340                                               def memory_usage(self, deep=False):
  ...
  1363         1         56.0     56.0      0.2          if hasattr(self.array, "memory_usage"):
  1364                                                       return self.array.memory_usage(deep=deep)
  1365
  1366         1         11.0     11.0      0.0          v = self.array.nbytes
  1367         1         18.0     18.0      0.1          if deep and is_object_dtype(self) and not PYPY:
  1368         1      34233.0  34233.0     99.7              v += lib.memory_usage_of_objects(self.array)
  1369         1          1.0      1.0      0.0          return v

THe % time column clearly points to lib.memory_usage_of_objects. This is a Cython function, so we can’t use line-profiler on it. But we know from the snakeviz output above that we eventually get to PandasArray.__getitem__

In [11]: %lprun -f pd.arrays.PandasArray.__getitem__ df.memory_usage(deep=True)
Timer unit: 1e-06 s

Total time: 0.041508 s
File: /Users/taugspurger/miniconda3/envs/pandas=1.0.1/lib/python3.8/site-packages/pandas/core/arrays/numpy_.py
Function: __getitem__ at line 232

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
   232                                               def __getitem__(self, item):
   233     10000       4246.0      0.4     10.2          if isinstance(item, type(self)):
   234                                                       item = item._ndarray
   235
   236     10000      25475.0      2.5     61.4          item = check_array_indexer(self, item)
   237
   238     10000       4394.0      0.4     10.6          result = self._ndarray[item]
   239     10000       4386.0      0.4     10.6          if not lib.is_scalar(item):
   240                                                       result = type(self)(result)
   241     10000       3007.0      0.3      7.2          return result

In this particular example, the most notable thing is that fact that we’re calling this function 10,000 times, which amounts to once per item on our 10,000 row DataFrame. Again, the details of this specific example and the fix aren’t too important, but the solution was to just stop doing that².

The fix was provided by @neilkg soon after the issue was identified, and crucially included a new asv benchmark for memory_usage with object dtypes. Hopefully we won’t regress on this again in the future.

Workflow issues

This setup is certainly better than nothing. But there are a few notable problems, some general and some specific to pandas:

Writing benchmarks is hard work (just like tests). There’s the general issue of writing and maintaining code. And on top of that, writing a good ASV benchmark requires some knowledge specific to ASV. And again, just like tests, your benchmarks can be trusted only as far as their coverage. For a large codebase like pandas you’ll need a decently large benchmark suite.

But that large benchmark suite comes with it’s own costs. Currently pandas’ full suite takes about 2 hours to run. This rules out running the benchmarks on most public CI providers. And even if we could finish it in time, we couldn’t really trust the results. These benchmarks, at least as written, really do need dedicated hardware to be stable over time. Pandas has a machine in my basement, but maintaining that has been a time-consuming, challenging process.

This is my current setup, which stuffs the benchmark server (the black Intel NUC) and a router next to my wife’s art storage. We reached this solution after my 2 year old unplugged the old setup (on my office floor) one too many times. Apologies for the poor cabling.

We deploy the benchmarks (for pandas and a few other NumFOCUS projects) using Ansible. The scripts get the benchmarks in place, Airflow to run them nightly, and supervisord to kick everything off. The outputs are rsynced over to the pandas webserver and served at https://pandas.pydata.org/speed/. You can see pandas’ at https://pandas.pydata.org/speed/pandas/. If this seems like a house of cards waiting to tumble, that’s because it is.

Pandas has applied for a NumFOCUS small development grant to improve our benchmark process. Ideally maintainers would be able to ask a bot @asv-bot run -b memory_usage which would kick off a process that pulled down the pull request and ran the requested benchmarks on a dedicated machine (that isn’t easily accessible by my children).

Recap

To summarize:

1. We need benchmarks to monitor performance, especially over time
2. We use tools like asv to organize and benchmark continuously
3. When regressions occur, we use snakeviz and line-profiler to diagnose the problem

Compatibility Code

Most libraries with dependencies will want to support multiple versions of that dependency. But supporting old version is a pain: it requires compatibility code, code that is around solely to get the same output from versions of a library. This post gives some advice on writing compatibility code.

1. Don’t write your own version parser
2. Centralize all version parsing
3. Use consistent version comparisons
4. Use Python’s argument unpacking
5. Clean up unused compatibility code

1. Don’t write your own version parser

It can be tempting just do something like

if pandas.__version__.split(".")[1] >= "25":
    ...

But that’s probably going to break, sometimes in unexpected ways. Use either distutils.version.LooseVersion or packaging.version.parse which handles all the edge cases.

PANDAS_VERSION = LooseVersion(pandas.__version__)

2. Centralize all version parsing in a _compat.py file

The first section of compatibility code is typically a version check. It can be tempting to do the version-check inline with the compatibility code

if LooseVersion(pandas.__version__) >= "0.25.0":
    return pandas.concat(args, sort=False)
else:
    return pandas.concat(args)

Rather than that, I recommend centralizing the version checks in a central _compat.py file that defines constants for each library version you need compatibility code for.

# library/_compat.py
import pandas


PANDAS_VERSION = LooseVersion(pandas.__version__)
PANDAS_0240 = PANDAS_VERSION >= "0.24.0
PANDAS_0250 = PANDAS_VERSION >= "0.25.0

This, combined with item 3, will make it easier to clean up your code (see below).

3. Use consistent version comparisons

Notice that I defined constants for each pandas version, PANDAS_0240, PANDAS_0250. Those mean “the installed version of pandas is at least this version”, since I used the >= comparison. You could instead define constants like

PANDAS_LT_0240 = PANDAS_VERSION < "0.24.0"

That works too, just ensure that you’re consistent.

4. Use Python’s argument unpacking

Python’s argument unpacking helps avoid code duplication when the signature of a function changes.

    param_grid = {"estimator__alpha": [0.1, 10]}
    if SK_022:
        kwargs = {}
    else:
        kwargs = {"iid": False}
    gs = sklearn.model_selection.GridSearchCV(clf, param_grid, cv=3, **kwargs)

Using *args, and **kwargs to pass through version-dependent arguments lets you have just a single call to the callable when the only difference is the arguments passed.

5. Clean up unused compatibility code

Actively developed libraries may eventually drop support for old versions of dependency libraries. At a minimum, this involves removing the old version from your test matrix and bumping your required version in your dependency list. But ideally you would also clean up the now-unused compatibility code. The strategies laid out here intend to make that as easy as possible.

Consider the following.

# library/core.py
import pandas
from ._comapt import PANDAS_0250


def f(args):
    ...

    if PANDAS_0250:
        return pandas.concat(args, sort=False)
    else:
        return pandas.concat(args)

Now suppose it’s the future and we want to drop support for pandas older than 0.25.x Now all the conditions checking if PANDAS_0250 are automatically true, so we’d

1. Delete PANDAS_0250 from _compat.py
2. Remove the import in core.py
3. Remove the if PANDAS_0250 check, and always have the True part of that condition

# library/core.py
import pandas

def f(args):
    ...
    return pandas.concat(args, sort=False)

I acknowledge that indirection can harm readability. In this case I think it’s warranted for actively maintained projects. Using inline version checks, perhaps with inconsistent comparisons, will make it harder to know when code is now unused. When integrated over the lifetime of the project, I find the strategies laid out here more readable.

Dask Summit Recap

Last week was the first Dask Developer Workshop. This brought together many of the core Dask developers and its heavy users to discuss the project. I want to share some of the experience with those who weren’t able to attend.

This was a great event. Aside from any technical discussions, it was ncie to meet all the people. From new acquaintences to people you’re on weekly calls with, it was great to interact with everyone.

The workshop

During our brief introductions, everyone included a one-phrase description of what they’d most-like to see improved in the project. These can roughly be grouped as

• Project health: more maintainers, more maintainer diversity, more commercial adoption
• Deployments: Support for heterogeneous clusters (e.g. some workers with different resources) on more cluster managers. Easier deployments for various use cases (single user vs. small team of scientists vs. enterprise IT managing things for a large team)
• Documentation: Including example
• Data Access: Loading data from various sources
• Reliability: Especially on adaptive clusters, as workers come and go.
• Features: Including things like approximate nearest neighbors, shared clients between futures, multi-column sorting, MultiIndex for dask.dataframe

One of the themes of the workshop was requests for honest, critical feedback about what needs to improve. Overall, people had great things to say about Dask and the various sub-projects but there’s always things to improve.

Dask sits at a pretty interesting place in the scientific Python ecosystem. It (and its users) are power-users of many libraries. It acts as a nice coordination point for many projects. We had maintainers from projects like NumPy, pandas, scikit-learn, Apache Arrow, cuDF, and others.

This post describes the start of a journey to get pandas’ documentation running on Binder. The end result is this nice button:

For a while now I’ve been jealous of Dask’s examples repository. That’s a repository containing a collection of Jupyter notebooks demonstrating Dask in action. It stitches together some tools to present a set of documentation that is both viewable as a static site at examples.dask.org, and as a executable notebooks on mybinder.

A bit of background on binder: it’s a tool for creating a shareable computing environment. This is perfect for introductory documentation. A prospective user may want to just try out a library to get a feel for it before they commit to installing. BinderHub is a tool for deploying binder services. You point a binderhub deployment (like mybinder) at a git repository with a collection of notebooks and an environment specification, and out comes your executable documentation.

Thanks to a lot of hard work by contributors and maintainers, the code examples in pandas’ documentation are already runnable (and this is verified on each commit). We use the IPython Sphinx Extension to execute examples and include their output. We write documentation like

.. ipython:: python

   import pandas as pd
   s = pd.Series([1, 2])
   s

Which is then executed and rendered in the HTML docs as

In [1]: import pandas as pd

In [2]: s = pd.Series([1, 2, 3])

In [3]: s
Out[3]:
0    1
1    2
2    3
dtype: int64

So we have the most important thing: a rich source of documentation that’s already runnable.

There were a couple barriers to just pointing binder at https://github.com/pandas-dev/pandas, however. First, binder builds on top of a tool called repo2docker. This is what takes your Git repository and turns it into a Docker image that users will be dropped into. When someone visits the URL, binder will first check to see if it’s built a docker image. If it’s already cached, then that will just be loaded. If not, binder will have to clone the repository and build it from scratch, a time-consuming process. Pandas receives 5-10 commits per day, meaning many users would visit the site and be stuck waiting for a 5-10 minute docker build.¹

Second, pandas uses Sphinx and ReST for its documentation. Binder needs a collection of Notebooks. Fortunately, the fine folks at QuantEcon (a fellow NumFOCUS project) wrote sphinxcontrib-jupyter, a tool for turning ReST files to Jupyter notebooks. Just what we needed.

So we had some great documentation that already runs, and a tool for converting ReST files to Jupyter notebooks. All the pieces were falling into place!

Unfortunately, my first attempt failed. sphinxcontrib-jupyter looks for directives like

.. code:: python

while pandas uses

.. ipython:: python

I started slogging down a path to teach sphinxcontrib-jupyter how to recognize the IPython directive pandas uses when my kid woke up from his nap. Feeling dejected I gave up.

But later in the day, I had the (obvious in hindsight) realization that we have plenty of tools for substituting lines of text. A few (non-obvious) lines of bash later and we were ready to go. All the .. ipython:: python directives were now .. code:: python. Moral of the story: take breaks.

My work currently lives in this repository, and the notebooks are runnable on mybinder. But the short version is

1. We include github.com/pandas-dev/pandas as a submodule (which repo2docker supports just fine)
2. We patch pandas Sphinx config to include sphinxcontrib-jupyter and its configuration
3. We patch pandas source docs to change the ipython directives to be .. code:: python directives.

I’m reasonably happy with how things are shaping up. I plan to migrate my repository to the pandas organization and propose a few changes to the pandas documentation (like a small header pointing from the rendered HTML docs to the binder). If you’d like to follow along, subscribe to this pandas issue.

I’m also hopeful that other projects can apply a similar approach to their documentation too.

This post describes a few protocols taking shape in the scientific Python community. On their own, each is powerful. Together, I think they enable for an explosion of creativity in the community.

Each of the protocols / interfaces we’ll consider deal with extending.

• NEP-13: NumPy __array_ufunc__
• NEP-18: NumPy __array_function__
• Pandas Extension types
• Custom Dask Collections

First, a bit of brief background on each.

NEP-13 and NEP-18, each deal with using the NumPy API on non-NumPy ndarray objects. For example, you might want to apply a ufunc like np.log to a Dask array.

>>> a = da.random.random((10, 10))
>>> np.log(a)
dask.array<log, shape=(10, 10), dtype=float64, chunksize=(10, 10)>

Prior to NEP-13, dask.array needed it’s own namespace of ufuncs like da.log, since np.log would convert the Dask array to an in-memory NumPy array (probably blowing up your machine’s memory). With __array_ufunc__ library authors and users can all just use NumPy ufuncs, without worrying about the type of the Array object.

While NEP-13 is limited to ufuncs, NEP-18 applies the same idea to most of the NumPy API. With NEP-18, libraries written to deal with NumPy ndarrays may suddenly support any object implementing __array_function__.

I highly recommend reading this blog post for more on the motivation for __array_function__. Ralph Gommers gave a nice talk on the current state of things at PyData Amsterdam 2019, though this is an active area of development.

Pandas added extension types to allow third-party libraries to solve domain-specific problems in a way that gels nicely with the rest of pandas. For example, cyberpandas handles network data, while geopandas handles geographic data. When both implement extension arrays it’s possible to operate on a dataset with a mixture of geographic and network data in the same DataFrame.

Finally, Dask defines a Collections Interface so that any object can be a first-class citizen within Dask. This is what ensures XArray’s DataArray and Dataset objects work well with Dask.

Series.__array_ufunc__

Now, onto the fun stuff: combining these interfaces across objects and libraries. https://github.com/pandas-dev/pandas/pull/23293 is a pull request adding Series.__array_ufunc__. There are a few subtleties, but the basic idea is that a ufunc applied to a Series should

1. Unbox the array (ndarray or extension array) from the Series
2. Apply the ufunc to the Series (honoring the array’s __array_ufunc__ if needed)
3. Rebox the output in a Series (with the original index and name)

For example, pandas’ SparseArray implements __array_ufunc__. It works by calling the ufunc twice, once on the sparse values (e.g. the non-zero values), and once on the scalar fill_value. The result is a new SparseArray with the same memory usage. With that PR, we achieve the same thing when operating on a Series containing an ExtensionArray.

>>> ser = pd.Series(pd.SparseArray([-10, 0, 10] + [0] * 100000))
>>> ser
0        -10
1          0
2         10
3          0
4          0
          ..
99998      0
99999      0
100000     0
100001     0
100002     0
Length: 100003, dtype: Sparse[int64, 0]

>>> n [20]: np.sign(ser)
0        -1
1         0
2         1
3         0
4         0
         ..
99998     0
99999     0
100000    0
100001    0
100002    0
Length: 100003, dtype: Sparse[int64, 0]

Previously, that would have converted the SparseArray to a dense NumPy array, blowing up your memory, slowing things down, and giving the incorrect result.

IPArray.__array_function__

To demonstrate __array_function__, we’ll implement it on IPArray.

    def __array_function__(self, func, types, args, kwargs):
        cls = type(self)
        if not all(issubclass(t, cls) for t in types):
            return NotImplemented
        return HANDLED_FUNCTIONS[func](*args, **kwargs)

IPArray is pretty domain-specific, so we place ourself down at the bottom priority by returning NotImplemented if there are any types we don’t recognize (we might consider handling Python’s stdlib ipaddres.IPv4Address and ipaddres.IPv6Address objects too).

And then we start implementing the interface. For example, concatenate.

@implements(np.concatenate)
def concatenate(arrays, axis=0, out=None):
    if axis != 0:
        raise NotImplementedError(f"Axis != 0 is not supported. (Got {axis}).")
    return IPArray(np.concatenate([array.data for array in arrays]))

With this, we can successfully concatenate two IPArrays

>>> a = cyberpandas.ip_range(4)
>>> b = cyberpandas.ip_range(10, 14)
>>> np.concatenate([a, b])
IPArray(['0.0.0.0', '0.0.0.1', '0.0.0.2', '0.0.0.3', '0.0.0.10', '0.0.0.11', '0.0.0.12', '0.0.0.13'])

Extending Dask

Finally, we may wish to make IPArray work well with dask.dataframe, to do normal cyberpandas operations in parallel, possibly distributed on a cluster. This requires a few changes:

1. Updating IPArray to work on either NumPy or Dask arrays
2. Implementing the Dask Collections interface on IPArray.
3. Registering an ip accessor with dask.dataframe, just like with pandas.

This is demonstrated in https://github.com/ContinuumIO/cyberpandas/pull/39

In [28]: ddf
Out[28]:
Dask DataFrame Structure:
                 A
npartitions=2
0               ip
6              ...
11             ...
Dask Name: from_pandas, 2 tasks

In [29]: ddf.A.ip.netmask()
Out[29]:
Dask Series Structure:
npartitions=2
0      ip
6     ...
11    ...
Name: A, dtype: ip
Dask Name: from-delayed, 22 tasks

In [30]: ddf.A.ip.netmask().compute()
Out[30]:
0     255.255.255.255
1     255.255.255.255
2     255.255.255.255
3     255.255.255.255
4     255.255.255.255
5     255.255.255.255
6     255.255.255.255
7     255.255.255.255
8     255.255.255.255
9     255.255.255.255
10    255.255.255.255
11    255.255.255.255
dtype: ip

Conclusion

I think that these points of extension.

Scikit-Learn 0.20.0 will contain some nice new features for working with tabular data. This blogpost will introduce those improvements with a small demo. We’ll then see how Dask-ML was able to piggyback on the work done by scikit-learn to offer a version that works well with Dask Arrays and DataFrames.

import dask
import dask.array as da
import dask.dataframe as dd
import numpy as np
import pandas as pd
import seaborn as sns
import fastparquet
from distributed import Client
from distributed.utils import format_bytes

Background

For the most part, Scikit-Learn uses NumPy ndarrays or SciPy sparse matricies for its in-memory data structures. This is great for many reasons, but one major drawback is that you can’t store heterogenous (AKA tabular) data in these containers. These are datasets where different columns of the table have different data types (some ints, some floats, some strings, etc.).

Pandas was built to work with tabular data. Scikit-Learn was built to work with NumPy ndarrays and SciPy sparse matricies. So there’s some friction when you use the two together. Perhaps someday things will be perfectly smooth, but it’s a challenging problem that will require work from several communities to fix. In this PyData Chicago talk, I discuss the differences between the two data models of scikit-learn and pandas, and some ways of working through it. The second half of the talk is mostly irrelevant now that ColumnTransformer is in scikit-learn.

ColumnTransformer in Scikit-Learn

At SciPy 2018, Joris Van den Bossche (a scikit-learn and pandas core developer) gives an update on some recent improvements to scikit-learn to make using pandas and scikit-learn together better.

The biggest addition is sklearn.compose.ColumnTransformer, a transformer for working with tabular data. The basic idea is to specify pairs of (column_selection, transformer). The transformer will be applied just to the selected columns, and the remaining columns can be passed through or dropped. Column selections can be integer positions (for arrays), names (for DataFrames) or a callable.

Here’s a small example on the “tips” dataset.

df = sns.load_dataset('tips')
df.head()

Our target is whether the tip was larger than 15%.

X = df.drop("tip", axis='columns')
y = df.tip / df.total_bill > 0.15

We’ll make a small pipeline that one-hot encodes the categorical columns (sex, smoker, day, time) before fitting a random forest. The numeric columns (total_bill, size) will be passed through as-is.

import sklearn.compose
import sklearn.ensemble
import sklearn.pipeline
import sklearn.preprocessing

We use make_column_transformer to create the ColumnTransformer.

categorical_columns = ['sex', 'smoker', 'day', 'time']
categorical_encoder = sklearn.preprocessing.OneHotEncoder(sparse=False)

transformers = sklearn.compose.make_column_transformer(
    (categorical_columns, categorical_encoder),
    remainder='passthrough'
)

This is just a regular scikit-learn estimator, which can be placed in a pipeline.

pipe = sklearn.pipeline.make_pipeline(
    transformers,
    sklearn.ensemble.RandomForestClassifier(n_estimators=100)
)

pipe.fit(X, y)
pipe.score(X, y)

1.0
    

We’ve likely overfitted, but that’s not really the point of this article. We’re more interested in the pre-processing side of things.

ColumnTransformer in Dask-ML

ColumnTransfomrer was added to Dask-ML in https://github.com/dask/dask-ml/pull/315. Ideally, we wouldn’t need that PR at all. We would prefer for dask’s collections (and pandas dataframes) to just be handled gracefully by scikit-learn. The main blocking issue is that the Python community doesn’t currently have a way to write “concatenate this list of array-like objects together” in a generic way. That’s being worked on in NEP-18.

So for now, if you want to use ColumnTransformer with dask objects, you’ll have to use dask_ml.compose.ColumnTransformer, otherwise your large Dask Array or DataFrame would be converted to an in-memory NumPy array.

As a footnote to this section, the initial PR in Dask-ML was much longer. I only needed to override one thing (the function _hstack used to glue the results back together). But that was being called from several places, and so I had to override all those places as well. I was able to work with the scikit-learn developers to make _hstack a staticmethod on ColumnTranformer, so any library wishing to extend ColumnTransformer can do so more easily now. The Dask project values working with the existing community.

Challenges with Scaling

Many strategies for dealing with large datasets rely on processing the data in chunks. That’s the basic idea behind Dask DataFrame: a Dask DataFrame consists of many pandas DataFrames. When you write ddf.column.value_counts(), Dask builds a task graph with many pandas.value_counts, and a final aggregation step so that you end up with the same end result.

But chunking can cause issues when there are variations in your dataset and the operation you’re applying depends on the data. For example, consider scikit-learn’s OneHotEncoder. By default, it looks at the data and creates a column for each unique value.

enc = sklearn.preprocessing.OneHotEncoder(sparse=False)
enc.fit_transform([['a'], ['a'], ['b'], ['c']])

array([[1., 0., 0.],
       [1., 0., 0.],
       [0., 1., 0.],
       [0., 0., 1.]])

But let’s suppose we wanted to process that in chunks of two, first [['a'], ['a']], then [['b'], ['c']].

enc.fit_transform([['a'], ['a']])

array([[1.],
       [1.]])

enc.fit_transform([['b'], ['c']])

array([[1., 0.],
       [0., 1.]])

We have a problem! Two in fact:

1. The shapes don’t match. The first batch only saw “a”, so the output shape is (2, 1). We can’t concatenate these results vertically
2. The meaning of the first column of the output has changed. In the first batch, the first column meant “a” was present. In the second batch, it meant “b” was present.

If we happened to know the set of possible values ahead of time, we could pass those to CategoricalEncoder. But storing that set of possible values separate from the data is fragile. It’d be better to store the possible values in the data type itself.

That’s exactly what pandas Categorical does. We can confidently know the number of columns in the categorical-encoded data by just looking at the type. Because this is so important in a distributed dataset context, dask_ml.preprocessing.OneHotEncoder differs from scikit-learn when passed categorical data: we use pandas’ categorical information.

A larger Example

We’ll work with the Criteo dataset. This has a mixture of numeric and categorical features. It’s also a large dataset, which presents some challenges for many pre-processing methods.

The full dataset is from http://labs.criteo.com/2013/12/download-terabyte-click-logs/. We’ll work with a sample.

client = Client()

ordinal_columns = [
    'category_0', 'category_1', 'category_2', 'category_3',
    'category_4', 'category_6', 'category_7', 'category_9',
    'category_10', 'category_11', 'category_13', 'category_14',
    'category_17', 'category_19', 'category_20', 'category_21',
    'category_22', 'category_23',
]

onehot_columns = [
    'category_5', 'category_8', 'category_12',
    'category_15', 'category_16', 'category_18',
    'category_24', 'category_25',
]

numeric_columns = [f'numeric_{i}' for i in range(13)]
columns = ['click'] + numeric_columns + onehot_columns + ordinal_columns

The raw data is a single large CSV. That’s been split with this script and I took a 10% sample with this script, which was written to a directory of parquet files. That’s what we’ll work with.

sample = dd.read_parquet("data/sample-10.parquet/")

# Convert unknown categorical to known.
# See note later on.

pf = fastparquet.ParquetFile("data/sample-10.parquet/part.0.parquet")
cats = pf.grab_cats(onehot_columns)

sample = sample.assign(**{
    col: sample[col].cat.set_categories(cats[col]) for col in onehot_columns
})

Our goal is to predict ‘click’ using the other columns.

y = sample['click']
X = sample.drop("click", axis='columns')

Now, let’s lay out our pre-processing pipeline. We have three types of columns

1. Numeric columns
2. Low-cardinality categorical columns
3. High-cardinality categorical columns

Each of those will be processed differently.

1. Numeric columns will have missing values filled with the column average and standard scaled
2. Low-cardinality categorical columns will be one-hot encoded
3. High-cardinality categorical columns will be deterministically hashed and standard scaled

You’ll probably want to quibble with some of these choices, but right now, I’m just interested in the ability to do these kinds of transformations at all.

We need to define a couple custom estimators, one for hashing the values of a dask dataframe, and one for converting a dask dataframe to a dask array.

import sklearn.base

def hash_block(x: pd.DataFrame) -> pd.DataFrame:
    """Hash the values in a DataFrame."""
    hashed = [
        pd.Series(pd.util.hash_array(data.values), index=x.index, name=col)
        for col, data in x.iteritems()
    ]
    return pd.concat(hashed, axis='columns')


class HashingEncoder(sklearn.base.BaseEstimator, sklearn.base.TransformerMixin):
    def fit(self, X, y=None):
        return self

    def transform(self, X, y=None):
        if isinstance(X, pd.DataFrame):
            return hash_block(X)
        elif isinstance(X, dd.DataFrame):
            return X.map_partitions(hash_block)
        else:
            raise ValueError("Unexpected type '{}' for 'X'".format(type(X)))


class ArrayConverter(sklearn.base.BaseEstimator, sklearn.base.TransformerMixin):
    """Convert a Dask DataFrame to a Dask Array with known lengths"""
    def __init__(self, lengths=None):
        self.lengths = lengths

    def fit(self, X, y=None):
        return self

    def transform(self, X, y=None):
        return X.to_dask_array(lengths=self.lengths)

For the final stage, Dask-ML needs to have a Dask Array with known chunk lengths. So let’s compute those ahead of time, and get a bit of info about how large the dataset is while we’re at it.

lengths = sample['click'].map_partitions(len)
nbytes = sample.memory_usage(deep=True).sum()

lengths, nbytes = dask.compute(lengths, nbytes)
lengths = tuple(lengths)

format_bytes(nbytes)

'19.20 GB'

We we’ll be working with about 20GB of data on a laptop with 16GB of RAM. We’ll clearly be relying on Dask to do the operations in parallel, while keeping things in a small memory footprint.

from dask_ml.compose import make_column_transformer
from dask_ml.preprocessing import StandardScaler, OneHotEncoder
from dask_ml.wrappers import Incremental
from dask_ml.impute import SimpleImputer

from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import FunctionTransformer
from sklearn.linear_model import SGDClassifier

Now for the pipeline.

onehot_encoder = OneHotEncoder(sparse=False)
hashing_encoder = HashingEncoder()
nan_imputer = SimpleImputer()

to_numeric = make_column_transformer(
    (onehot_columns, onehot_encoder),
    (ordinal_columns, hashing_encoder),
    remainder='passthrough',
)

fill_na = make_column_transformer(
    (numeric_columns, nan_imputer),
    remainder='passthrough'
)

scaler = make_column_transformer(
    (list(numeric_columns) + list(ordinal_columns), StandardScaler()),
    remainder='passthrough'
)


clf = Incremental(
    SGDClassifier(loss='log',
                  random_state=0,
                  max_iter=1000)
)

pipe = make_pipeline(to_numeric, fill_na, scaler, ArrayConverter(lengths=lengths), clf)
pipe

Pipeline(memory=None,
     steps=[('columntransformer-1', ColumnTransformer(n_jobs=1, preserve_dataframe=True, remainder='passthrough',
         sparse_threshold=0.3, transformer_weights=None,
         transformers=[('onehotencoder', OneHotEncoder(categorical_features=None, categories='auto',
       dtype=<class 'numpy.float6...ion=0.1, verbose=0, warm_start=False),
      random_state=None, scoring=None, shuffle_blocks=True))])

Overall it reads pretty similarly to how we described it in prose. We specify

1. Onehot the low-cardinality categoricals, hash the others
2. Fill missing values in the numeric columns
3. Standard scale the numeric and hashed columns
4. Fit the incremental SGD

And again, these ColumnTransformers are just estimators so we stick them in a regular scikit-learn Pipeline before calling .fit:

%%time pipe.fit(X, y.to_dask_array(lengths=lengths), incremental__classes=[0, 1])

CPU times: user 7min 7s, sys: 41.6 s, total: 7min 48s
Wall time: 16min 42s
Pipeline(memory=None,
steps=[(‘columntransformer-1’, ColumnTransformer(n_jobs=1, preserve_dataframe=True, remainder=‘passthrough’,
sparse_threshold=0.3, transformer_weights=None,
transformers=[(‘onehotencoder’, OneHotEncoder(categorical_features=None, categories=‘auto’,
dtype=<class ’numpy.float6…ion=0.1, verbose=0, warm_start=False),
random_state=None, scoring=None, shuffle_blocks=True))])

Discussion

Some aspects of this workflow could be improved.

1. Dask, fastparquet, pyarrow, and pandas don’t currently have a way to specify the categorical dtype of a column split across many files. Each file (parition) is treated independently. This results in categorials with unknown categories in the Dask DataFrame. Since we know that the categories are all the same, we’re able to read in the first files categories and assign those to the entire DataFrame. But this is a bit fragile, as it relies on an assumption not necessarily guaranteed by the file structure.

2. There’s of IO. As written, each stage of the pipeline that has to see the data does a full read of the dataset. We end up reading the entire dataset something like 5 times. https://github.com/dask/dask-ml/issues/192 has some discussion on ways we can progress through a pipeline. If your pipeline consists entirely of estimators that learn incrementally, it may make sense to send each block of data through the entire pipeline, rather than sending all the data to the first step, then all the data to the second, and so on. I’ll note, however, that you can avoid the redundant IO by loading your data into distributed RAM on a Dask cluster. But I was just trying things out on my laptop.

Still, it’s worth noting that we’ve successfully fit a reasonably complex pipeline on a larger-than-RAM dataset using our laptop. That’s something!

ColumnTransformer will be available in scikit-learn 0.20.0. This also contains the changes for distributed joblib I blogged about earlier. The first release candidate is available now.

For more, visit the Dask, Dask-ML, and scikit-learn documentation.

This work is supported by Anaconda Inc.

This post describes a recent improvement made to TPOT. TPOT is an automated machine learning library for Python. It does some feature engineering and hyper-parameter optimization for you. TPOT uses genetic algorithms to evaluate which models are performing well and how to choose new models to try out in the next generation.

Parallelizing TPOT

In TPOT-730, we made some modifications to TPOT to support distributed training. As a TPOT user, the only changes you need to make to your code are

1. Connect a client to your Dask Cluster
2. Specify the use_dask=True argument to your TPOT estimator

From there, all the training will use your cluster of machines. This screencast shows an example on an 80-core Dask cluster.

Commentary

Fitting a TPOT estimator consists of several stages. The bulk of the time is spent evaluating individual scikit-learn pipelines. Dask-ML already had code for splitting apart a scikit-learn Pipeline.fit call into individual tasks. This is used in Dask-ML’s hyper-parameter optimization to avoid repeating work. We were able to drop-in Dask-ML’s fit and scoring method for the one already used in TPOT. That small change allows fitting the many individual models in a generation to be done on a cluster.

There’s still some room for improvement. Internal to TPOT, some time is spent determining the next set of models to try out (this is the “mutation and crossover phase”). That’s not (yet) been parallelized with Dask, so you’ll notice some periods of inactivity on the cluster.

Next Steps

This will be available in the next release of TPOT. You can try out a small example now on the dask-examples binder.

Stepping back a bit, I think this is a good example of how libraries can use Dask internally to parallelize workloads for their users. Deep down in TPOT there was a single method for fitting many scikit-learn models on some data and collecting the results. Dask-ML has code for building a task graph that does the same thing. We were able to swap out the eager TPOT code for the lazy dask version, and get things distributed on a cluster. Projects like xarray have been able to do a similar thing with dask Arrays in place of NumPy arrays. If Dask-ML hadn’t already had that code, dask.delayed could have been used instead.

If you have a library that you think could take advantage of Dask, please reach out!

The other day, I put up a Twitter poll asking a simple question: What’s the type of series.values?

Pop Quiz! What are the possible results for the following:

>>> type(pandas.Series.values)

— Tom Augspurger (@TomAugspurger) August 6, 2018

I was a bit limited for space, so I’ll expand on the options here. Choose as many as you want.

1. NumPy ndarray
2. pandas Categorical (or all of the above)
3. An Index or any of it’s subclasses (DatetimeIndex, CategoricalIndex, RangeIndex, etc.) (or all of the above)
4. None or all of the above

I was prompted to write this post because a.) this is an (unfortunately) confusing topic and b.) it’s undergoing a lot of change right now (and, c.) I had this awesome title in my head).

The Answer

Unfortunately I kind of messed up the poll. Things are even more complex than I initially thought.

As best I can tell, the possible types for series.values are

• NumPy ndarray
• pandas Categorical
• pandas SparseArray (I forgot about this one in the poll)

So going with the cop-out “best-available” answer, I would have said that 2 was the best answer in the poll. SparseArray is technically and ndarray subclass (for now), so technically 2 is correct, but that’s a few too many technicallys for me.

The Explanation

So, that’s a bit of a mess. How’d we get here? Or, stepping back a bit, what even is an array? What’s a dataframe?

NumPy arrays are N-dimensional and homogenous. Every element in the array has to have the same data type.

Pandas dataframes are 2-dimensional and heterogenous. Different columns can have different data types. But every element in a single column (Series) has the same data type. I like to think of DataFrames as containers for Series. Stepping down a dimension, I think of Series as containers for 1-D arrays. In an ideal world, we could say Series are containers for NumPy ararys, but that’s not quite the case.

While there’s a lot of overlap between the pandas and NumPy communites, there are still differences. Pandas users place different value on different features, so pandas has restricted and extended NumPy’s type system in a few directions. For example, early Pandas users (many of them in the financial sector) needed datetimes with timezones, but didn’t really care about lower-precision timestamps like datetime64[D]. So pandas limited its scope to just nanosecond-precision datetimes (datetime64[ns]) and extended it with some metedata for the timezone. Likewise for Categorical, period, sparse, interval, etc.

So back to Series.values; pandas had a choice: should Series.values always be a NumPy array, even if it means losing information like the timezone or categories, and even if it’s slow or could exhaust your memory (large categorical or sparse arrays)? Or should it faithfully represent the data, even if that means not returning an ndarray?

I don’t think there’s a clear answer to this question. Both options have their downsides. In the end, we ended up with a messy compromise, where some things return ndarrays, some things return something else (Categorical), and some things do a bit of conversion before returning an ndarary.

For example, off the top of your head, do you know what the type of Series.values is for timezone-aware data?

In [2]: pd.Series(pd.date_range('2017', periods=4, tz='US/Eastern'))
Out[2]:
0   2017-01-01 00:00:00-05:00
1   2017-01-02 00:00:00-05:00
2   2017-01-03 00:00:00-05:00
3   2017-01-04 00:00:00-05:00
dtype: datetime64[ns, US/Eastern]

In [3]: pd.Series(pd.date_range('2017', periods=4, tz='US/Eastern')).values
Out[3]:
array(['2017-01-01T05:00:00.000000000', '2017-01-02T05:00:00.000000000',
       '2017-01-03T05:00:00.000000000', '2017-01-04T05:00:00.000000000'],
      dtype='datetime64[ns]')

With the wisdom of Solomon, we decided to have it both ways; the values are converted to UTC and the timezone is dropped. I don’t think anyone would claim this is ideal, but it was backwards compatibile-ish. Given the constraints, it wasn’t the worst choice in the world.

The Near Future

In pandas 0.24, we’ll (hopefully) have a good answer for what series.values is: a NumPy array or an ExtensionArray. For regular data types represented by NumPy, you’ll get an ndarray. For extension types (implemented in pandas or elsewhere) you’ll get an ExtensionArray. If you’re using Series.values, you can rely on the set of methods common to each.

But that raises the question: why are you using .values in the first place? There are some legitmate use cases (disabling automatic alignment, for example), but for many things, passing a Series will hopefully work as well as a NumPy array. To users of pandas, I recommend avoiding .values as much as possible. If you know that you need an ndarray, you’re probably best of using np.asarray(series). That will do the right thing for any data type.

The Far Future

I’m hopeful that some day all we’ll have a common language for these data types. There’s a lot going on in the numeric Python ecosystem right now. Stay tuned!

This is part 1 in my series on writing modern idiomatic pandas.

• Modern Pandas
• Method Chaining
• Indexes
• Fast Pandas
• Tidy Data
• Visualization
• Time Series
• Scaling

As I sit down to write this, the third-most popular pandas question on StackOverflow covers how to use pandas for large datasets. This is in tension with the fact that a pandas DataFrame is an in memory container. You can’t have a DataFrame larger than your machine’s RAM. In practice, your available RAM should be several times the size of your dataset, as you or pandas will have to make intermediate copies as part of the analysis.

Historically, pandas users have scaled to larger datasets by switching away from pandas or using iteration. Both of these are perfectly valid approaches, but changing your workflow in response to scaling data is unfortunate. I use pandas because it’s a pleasant experience, and I would like that experience to scale to larger datasets. That’s what Dask, a parallel computing library, enables. We’ll discuss Dask in detail later. But first, let’s work through scaling a simple analysis to a larger than memory dataset.

Our task is to find the 100 most-common occupations reported in the FEC’s individual contributions dataest. The files are split by election cycle (2007-2008, 2009-2010, …). You can find some scripts for downloading the data in this repository. My laptop can read in each cycle’s file individually, but the full dataset is too large to read in at once. Let’s read in just 2010’s file, and do the “small data” version.

from pathlib import Path

import pandas as pd
import seaborn as sns

df = pd.read_parquet("data/indiv-10.parq", columns=['occupation'], engine='pyarrow')

most_common = df.occupation.value_counts().nlargest(100)
most_common

    RETIRED                    279775
    ATTORNEY                   166768
    PRESIDENT                   81336
    PHYSICIAN                   73015
    HOMEMAKER                   66057
                                ...  
    C.E.O.                       1945
    EMERGENCY PHYSICIAN          1944
    BUSINESS EXECUTIVE           1924
    BUSINESS REPRESENTATIVE      1879
    GOVERNMENT AFFAIRS           1867
    Name: occupation, Length: 100, dtype: int64

After reading in the file, our actual analysis is a simple 1-liner using two operations built into pandas. Truly, the best of all possible worlds.

Next, we’ll do the analysis for the entire dataset, which is larger than memory, in two ways. First we’ll use just pandas and iteration. Then we’ll use Dask.

Using Iteration

To do this with just pandas we have to rewrite our code, taking care to never have too much data in RAM at once. We will

1. Create a global total_counts Series that contains the counts from all of the files processed so far
2. Read in a file
3. Compute a temporary variable counts with the counts for just this file
4. Add that temporary counts into the global total_counts
5. Select the 100 largest with .nlargest

This works since the total_counts Series is relatively small, and each year’s data fits in RAM individually. Our peak memory usage should be the size of the largest individual cycle (2015-2016) plus the size of total_counts (which we can essentially ignore).

files = sorted(Path("data/").glob("indiv-*.parq"))

total_counts = pd.Series()

for year in files:
    df = pd.read_parquet(year, columns=['occupation'],
                         engine="pyarrow")
    counts = df.occupation.value_counts()
    total_counts = total_counts.add(counts, fill_value=0)

total_counts = total_counts.nlargest(100).sort_values(ascending=False)

RETIRED                    4769520
NOT EMPLOYED               2656988
ATTORNEY                   1340434
PHYSICIAN                   659082
HOMEMAKER                   494187
                            ...   
CHIEF EXECUTIVE OFFICER      26551
SURGEON                      25521
EDITOR                       25457
OPERATOR                     25151
ORTHOPAEDIC SURGEON          24384
Name: occupation, Length: 100, dtype: int64

While this works, our small one-liner has ballooned in size (and complexity; should you really have to know about Series.add’s fill_value parameter for this simple analysis?). If only there was a better way…

Using Dask

With Dask, we essentially recover our original code. We’ll change our import to use dask.dataframe.read_parquet, which returns a Dask DataFrame.

import dask.dataframe as dd

df = dd.read_parquet("data/indiv-*.parquet", engine='pyarrow', columns=['occupation'])

most_common = df.occupation.value_counts().nlargest(100)
most_common.compute().sort_values(ascending=False)

RETIRED                    4769520
NOT EMPLOYED               2656988
ATTORNEY                   1340434
PHYSICIAN                   659082
HOMEMAKER                   494187
                            ...   
CHIEF EXECUTIVE OFFICER      26551
SURGEON                      25521
EDITOR                       25457
OPERATOR                     25151
ORTHOPAEDIC SURGEON          24384
Name: occupation, Length: 100, dtype: int64

There are a couple differences from the original pandas version, which we’ll discuss next, but overall I hope you agree that the Dask version is nicer than the version using iteration.

Dask

Now that we’ve seen dask.dataframe in action, let’s step back and discuss Dask a bit. Dask is an open-source project that natively parallizes Python. I’m a happy user of and contributor to Dask.

At a high-level, Dask provides familiar APIs for large N-dimensional arrays, large DataFrames, and familiar ways to parallelize custom algorithms.

At a low-level, each of these is built on high-performance task scheduling that executes operations in parallel. The low-level details aren’t too important; all we care about is that

1. Dask works with task graphs (tasks: functions to call on data, and graphs: the relationships between tasks).
2. This is a flexible and performant way to parallelize many different kinds of problems.

To understand point 1, let’s examine the difference between a Dask DataFrame and a pandas DataFrame. When we read in df with dd.read_parquet, we received a Dask DataFrame.

df

A Dask DataFrame consists of many pandas DataFrames arranged by the index. Dask is really just coordinating these pandas DataFrames.

All the actual computation (reading from disk, computing the value counts, etc.) eventually use pandas internally. If I do df.occupation.str.len, Dask will coordinate calling pandas.Series.str.len on each of the pandas DataFrames.

Those reading carefully will notice a problem with the statement “A Dask DataFrame consists of many pandas DataFrames”. Our initial problem was that we didn’t have enough memory for those DataFrames! How can Dask be coordinating DataFrames if there isn’t enough memory? This brings us to the second major difference: Dask DataFrames (and arrays) are lazy. Operations on them don’t execute and produce the final result immediately. Rather, calling methods on them builds up a task graph.

We can visualize task graphs using graphviz. For the blog, I’ve trimmed down the example to be a subset of the entire graph.

df.visualize(rankdir='LR')

df (the dask DataFrame consisting of many pandas DataFrames) has a task graph with 5 calls to a parquet reader (one for each file), each of which produces a DataFrame when called.

Calling additional methods on df adds additional tasks to this graph. For example, our most_common Series has three additional calls

• Select the occupation column (__getitem__)
• Perform the value counts
• Select the 100 largest values

most_common = df.occupation.value_counts().nlargest(100)
most_common

    Dask Series Structure:
    npartitions=1
        int64
          ...
    Name: occupation, dtype: int64
    Dask Name: series-nlargest-agg, 113 tasks

Which we can visualize.

most_common.visualize(rankdir='LR')

So most_common doesn’t hold the actual answer yet. Instead, it holds a recipe for the answer; a list of all the steps to take to get the concrete result. One way to ask for the result is with the compute method.

most_common.compute()

    RETIRED                    4769520
    NOT EMPLOYED               2656988
    ATTORNEY                   1340434
    PHYSICIAN                   659082
    HOMEMAKER                   494187
                                ...   
    CHIEF EXECUTIVE OFFICER      26551
    SURGEON                      25521
    EDITOR                       25457
    OPERATOR                     25151
    ORTHOPAEDIC SURGEON          24384
    Name: occupation, Length: 100, dtype: int64

At this point, the task graph is handed to a scheduler, which is responsible for executing a task graph. Schedulers can analyze a task graph and find sections that can run in parallel. (Dask includes several schedulers. See the scheduling documentation for how to choose, though Dask has good defaults.)

So that’s a high-level tour of how Dask works:

1. Various collections collections like dask.dataframe and dask.array provide users familiar APIs for working with large datasets.
2. Computations are represented as a task graph. These graphs could be built by hand, or more commonly built by one of the collections.
3. Dask schedulers run task graphs in parallel (potentially distributed across a cluster), reusing libraries like NumPy and pandas to do the computations.

Let’s finish off this post by continuing to explore the FEC dataset with Dask. At this point, we’ll use the distributed scheduler for it’s nice diagnostics.

import dask.dataframe as dd
from dask import compute
from dask.distributed import Client
import seaborn as sns

client = Client(processes=False)

Calling Client without providing a scheduler address will make a local “cluster” of threads or processes on your machine. There are many ways to deploy a Dask cluster onto an actual cluster of machines, though we’re particularly fond of Kubernetes. This highlights one of my favorite features of Dask: it scales down to use a handful of threads on a laptop or up to a cluster with thousands of nodes. Dask can comfortably handle medium-sized datasets (dozens of GBs, so larger than RAM) on a laptop. Or it can scale up to very large datasets with a cluster.

individual_cols = ['cmte_id', 'entity_tp', 'employer', 'occupation',
                   'transaction_dt', 'transaction_amt']

indiv = dd.read_parquet('data/indiv-*.parq',
                        columns=individual_cols,
                        engine="pyarrow")
indiv

We can compute summary statistics like the average mean and standard deviation of the transaction amount:

avg_transaction = indiv.transaction_amt.mean()

We can answer questions like “Which employer’s employees donated the most?”

total_by_employee = (
    indiv.groupby('employer')
        .transaction_amt.sum()
        .nlargest(10)
)

Or “what is the average amount donated per occupation?”

avg_by_occupation = (
    indiv.groupby("occupation")
        .transaction_amt.mean()
        .nlargest(10)
)

Since Dask is lazy, we haven’t actually computed anything.

total_by_employee

    Dask Series Structure:
    npartitions=1
        int64
          ...
    Name: transaction_amt, dtype: int64
    Dask Name: series-nlargest-agg, 13 tasks

avg_transaction, avg_by_occupation and total_by_employee are three separate computations (they have different task graphs), but we know they share some structure: they’re all reading in the same data, they might select the same subset of columns, and so on. Dask is able to avoid redundant computation when you use the top-level dask.compute function.

%%time
avg_transaction, by_employee, by_occupation = compute(
    avg_transaction, total_by_employee, avg_by_occupation
)

    CPU times: user 57.5 s, sys: 14.4 s, total: 1min 11s
    Wall time: 54.9 s

avg_transaction

    566.0899206077507

by_employee

    employer
    RETIRED                1019973117
    SELF-EMPLOYED           834547641
    SELF                    537402882
    SELF EMPLOYED           447363032
    NONE                    418011322
    HOMEMAKER               355195126
    NOT EMPLOYED            345770418
    FAHR, LLC               166679844
    CANDIDATE                75186830
    ADELSON DRUG CLINIC      53358500
    Name: transaction_amt, dtype: int64

by_occupation

    occupation
    CHAIRMAN CEO & FOUNDER                   1,023,333.33
    PAULSON AND CO., INC.                    1,000,000.00
    CO-FOUNDING DIRECTOR                       875,000.00
    CHAIRMAN/CHIEF TECHNOLOGY OFFICER          750,350.00
    CO-FOUNDER, DIRECTOR, CHIEF INFORMATIO     675,000.00
    CO-FOUNDER, DIRECTOR                       550,933.33
    MOORE CAPITAL GROUP, LP                    500,000.00
    PERRY HOMES                                500,000.00
    OWNER, FOUNDER AND CEO                     500,000.00
    CHIEF EXECUTIVE OFFICER/PRODUCER           500,000.00
    Name: transaction_amt, dtype: float64

Things like filtering work well. Let’s find the 10 most common occupations and filter the dataset down to just those.

top_occupations = (
    indiv.occupation.value_counts()
        .nlargest(10).index
).compute()
top_occupations

    Index(['RETIRED', 'NOT EMPLOYED', 'ATTORNEY', 'PHYSICIAN', 'HOMEMAKER',
           'PRESIDENT', 'PROFESSOR', 'CONSULTANT', 'EXECUTIVE', 'ENGINEER'],
          dtype='object')

We’ll filter the raw records down to just the ones from those occupations. Then we’ll compute a few summary statistics on the transaction amounts for each group.

donations = (
    indiv[indiv.occupation.isin(top_occupations)]
        .groupby("occupation")
        .transaction_amt
        .agg(['count', 'mean', 'sum', 'max'])
)

total_avg, occupation_avg = compute(indiv.transaction_amt.mean(),
                                    donations['mean'])

These are small, concrete results so we can turn to familiar tools like matplotlib to visualize the result.

ax = occupation_avg.sort_values(ascending=False).plot.barh(color='k', width=0.9);
lim = ax.get_ylim()
ax.vlines(total_avg, *lim, color='C1', linewidth=3)
ax.legend(['Average donation'])
ax.set(xlabel="Donation Amount", title="Average Dontation by Occupation")
sns.despine()

Dask inherits all of pandas’ great time-series support. We can get the total amount donated per day using a resample.

daily = (
    indiv[['transaction_dt', 'transaction_amt']].dropna()
        .set_index('transaction_dt')['transaction_amt']
        .resample("D")
        .sum()
).compute()
daily

    1916-01-23    1000
    1916-01-24       0
    1916-01-25       0
    1916-01-26       0
    1916-01-27       0
                  ... 
    2201-05-29       0
    2201-05-30       0
    2201-05-31       0
    2201-06-01       0
    2201-06-02    2000
    Name: transaction_amt, Length: 104226, dtype: int64

It seems like we have some bad data. This should just be 2007-2016. We’ll filter it down to the real subset before plotting. Notice that the seamless transition from dask.dataframe operations above, to pandas operations below.

subset = daily.loc['2011':'2016']
ax = subset.div(1000).plot(figsize=(12, 6))
ax.set(ylim=0, title="Daily Donations", ylabel="$ (thousands)",)
sns.despine();

Joining

Like pandas, Dask supports joining together multiple datasets.

Individual donations are made to committees. Committees are what make the actual expenditures (buying a TV ad). Some committees are directly tied to a candidate (this are campaign committees). Other committees are tied to a group (like the Republican National Committee). Either may be tied to a party.

Let’s read in the committees. The total number of committees is small, so we’ll .compute immediately to get a pandas DataFrame (the reads still happen in parallel!).

committee_cols = ['cmte_id', 'cmte_nm', 'cmte_tp', 'cmte_pty_affiliation']
cm = dd.read_parquet("data/cm-*.parq",
                     columns=committee_cols).compute()

# Some committees change thier name, but the ID stays the same
cm = cm.groupby('cmte_id').last()
cm

We’ll use dd.merge, which is analogous to pd.merge for joining a Dask DataFrame with a pandas or Dask DataFrame.

indiv = indiv[(indiv.transaction_dt >= pd.Timestamp("2007-01-01")) &
              (indiv.transaction_dt <= pd.Timestamp("2018-01-01"))]

df2 = dd.merge(indiv, cm.reset_index(), on='cmte_id')
df2

Now we can find which party raised more over the course of each election. We’ll group by the day and party and sum the transaction amounts.

indiv = indiv.repartition(npartitions=10)
df2 = dd.merge(indiv, cm.reset_index(), on='cmte_id')
df2

party_donations = (
    df2.groupby([df2.transaction_dt, 'cmte_pty_affiliation'])
       .transaction_amt.sum()
).compute().sort_index()

We’ll filter that down to just Republican and Democrats and plot.

ax = (
    party_donations.loc[:, ['REP', 'DEM']]
        .unstack("cmte_pty_affiliation").iloc[1:-2]
        .rolling('30D').mean().plot(color=['C0', 'C3'], figsize=(12, 6),
                                    linewidth=3)
)
sns.despine()
ax.set(title="Daily Donations (30-D Moving Average)", xlabel="Date");

Try It Out!

So that’s a taste of Dask. Next time you hit a scaling problem with pandas (or NumPy, scikit-learn, or your custom code), feel free to

pip install dask[complete]

or

conda install dask

The dask homepage has links to all the relevant documentation, and binder notebooks where you can try out Dask before installing.

As always, reach out to me on Twitter or in the comments if you have anything to share.

This work is supported by Anaconda Inc and the Data Driven Discovery Initiative from the Moore Foundation.

dask-ml 0.4.1 was released today with a few enhancements. See the changelog for all the changes from 0.4.0.

Conda packages are available on conda-forge

$ conda install -c conda-forge dask-ml

and wheels and the source are available on PyPI

$ pip install dask-ml

I wanted to highlight one change, that touches on a topic I mentioned in my first post on scalable Machine Learning. I discussed how, in my limited experience, a common workflow was to train on a small batch of data and predict for a much larger set of data. The training data easily fits in memory on a single machine, but the full dataset does not.

A new meta-estimator, ParallelPostFit helps with this common case. It’s a meta-estimator that wraps a regular scikit-learn estimator, similar to how GridSearchCV wraps an estimator. The .fit method is very simple; it just calls the underlying estimator’s .fit method and copies over the learned attributes. This means ParalellPostFit is not suitable for training on large datasets. It is, however, perfect for post-fit tasks like .predict, or .transform.

As an example, we’ll fit a scikit-learn GradientBoostingClassifier on a small in-memory dataset.

>>> from sklearn.ensemble import GradientBoostingClassifier
>>> import sklearn.datasets
>>> import dask_ml.datasets

>>> X, y = sklearn.datasets.make_classification(n_samples=1000,
...                                             random_state=0)
>>> clf = ParallelPostFit(estimator=GradientBoostingClassifier())
>>> clf.fit(X, y)
ParallelPostFit(estimator=GradientBoostingClassifier(...))

Nothing special so far. But now, let’s suppose we had a “large” dataset for prediction. We’ll use dask_ml.datasets.make_classification, but in practice you would read this from a file system or database.

>>> X_big, y_big = dask_ml.datasets.make_classification(n_samples=100000,
                                                        chunks=1000,
                                                        random_state=0)

In this case we have a dataset with 100,000 samples split into blocks of 1,000. We can now predict for this large dataset.

>>> clf.predict(X)
dask.array<predict, shape=(10000,), dtype=int64, chunksize=(1000,)>

Now things are different. ParallelPostFit.predict, .predict_proba, and .transform, all return dask arrays instead of immediately computing the result. We’ve built up task graph of computations to be performed, which allows dask to step in and compute things in parallel. When you’re ready for the answer, call compute:

>>> clf.predict_proba(X).compute()
array([[0.99141094, 0.00858906],
       [0.93178389, 0.06821611],
       [0.99129105, 0.00870895],
       ...,
       [0.97996652, 0.02003348],
       [0.98087444, 0.01912556],
       [0.99407016, 0.00592984]])

At that point the dask scheduler comes in and executes your compute in parallel, using all the cores of your laptop or workstation, or all the machines on your cluster.

ParallelPostFit “fixes” a couple of issues in scikit-learn outside of scikit-learn itself

• parallel predict
• parallel transform

If you’re able to depend on dask and dask-ml, consider giving ParallelPostFit a shot and let me know how it turns out. For estimators whose predict is relatively expensive and not already parallelized, ParallelPostFit can give a nice performance boost.

Even if the underlying estimator’s predict or tranform method is cheap or parallelized, ParallelPostFit does still help with distributed the work on all the machines in your cluster, or doing the operation out-of-core.

Thanks to all the contributors who worked on this release.

This is a status update on some enhancements for pandas. The goal of the work is to store things that are sufficiently array-like in a pandas DataFrame, even if they aren’t a regular NumPy array. Pandas already does this in a few places for some blessed types (like Categorical); we’d like to open that up to anybody.

A couple months ago, a client came to Anaconda with a problem: they have a bunch of IP Address data that they’d like to work with in pandas. They didn’t just want to make a NumPy array of IP addresses for a few reasons:

1. IPv6 addresses are 128 bits, so they can’t use a specialized NumPy dtype. It would have to be an object array, which will be slow for their large datasets.
2. IP Addresses have special structure. They’d like to use this structure for special methods like is_reserved.
3. It’s much better to put the knowledge of types in the library, rather than relying on analysts to know that this column of objects or strings is actually this other special type.

I wrote up a proposal to gauge interest from the community for adding an IP Address dtype to pandas. The general sentiment was that an IP addresses were too specialized for inclusion pandas (which matched my own feelings). But, the community was interested in allowing 3rd party libraries to define their own types and having pandas “do the right thing” when it encounters them.

Pandas Internals

While not technically true, you could reasonably describe a DataFrame as a dictionary of NumPy arrays. There are a few complications that invalidate that caricature , but the one I want to focus on is pandas’ extension dtypes.

Pandas has extended NumPy’s type system in a few cases. For the most part, this involves tricking pandas.DataFrame and pandas.Series into thinking that the object passed to it is a single array, when in fact it’s multiple arrays, or an array plus a bit of extra metadata.

1. datetime64[ns] with a timezone. A regular numpy.datetime64[ns] array (which is really just an array of integers) plus some metadata for the timezone.
2. Period: An array of integer ordinals and some metadata about the frequency.
3. Categorical: two arrays: one with the unique set of categories and a second array of codes, the positions in categories.
4. Interval: Two arrays, one for the left-hand endpoints and one for the right-hand endpoints.

So our definition of a pandas.DataFrame is now “A dictionary of NumPy arrays, or one of pandas’ extension types.” Internal to pandas, we have checks for “is this thing an extension dtype? If so take this special path.” To the user, it looks like a Categorical is just a regular column, but internally, it’s a bit messier.

Anyway, the upshot of my proposal was to make changes to pandas' internals to support 3rd-party objects going down that “is this an extension dtype” path.

Pandas’ Array Interface

To support external libraries defining extension array types, we defined an interface.

In pandas-19268 we laid out exactly what pandas considers sufficiently “array-like” for an extension array type. When pandas comes across one of these array-like objects, it avoids the previous behavior of just storing the data in a NumPy array of objects. The interface includes things like

• What type of scalars do you hold?
• How do I convert you to a NumPy array?
• __getitem__

Most things should be pretty straightforward to implement. In the test suit, we have a 60-line implementation for storing decimal.Decimal objects in a Series.

It’s important to emphasize that pandas’ ExtensionArray is not another array implementation. It’s just an agreement between pandas and your library that your array-like object (which may be a NumPy array, many NumPy arrays, an Arrow array, a list, anything really) that satisfies the proper semantics for storage inside a Series or DataFrame.

With those changes, I’ve been able to prototype a small library (named… cyberpandas) for storing arrays of IP Addresses. It defines IPAddress, an array-like container for IP Addresses. For this blogpost, the only relevant implementation detail is that IP Addresses are stored as a NumPy structured array with two uint64 fields. So we’re making pandas treat this 2-D array as a single array, like how Interval works. Here’s a taste:

As a taste for what’s possible, here’s a preview of our IP Address library, cyberpandas.

In [1]: import cyberpandas

In [2]: import pandas as pd

In [3]: ips = cyberpandas.IPAddress([
   ...:     '0.0.0.0',
   ...:     '192.168.1.1',
   ...:     '2001:0db8:85a3:0000:0000:8a2e:0370:7334',
   ...: ])

In [4]: ips
Out[4]: IPAddress(['0.0.0.0', '192.168.1.1', '2001:db8:85a3::8a2e:370:7334'])

In [5]: ips.data
Out[5]:
array([(                  0,               0),
       (                  0,      3232235777),
       (2306139570357600256, 151930230829876)],
      dtype=[('hi', '>u8'), ('lo', '>u8')])

ips satisfies pandas’ ExtensionArray interface, so it can be stored inside pandas’ containers.

In [6]: ser = pd.Series(ips)

In [7]: ser
Out[7]:
0                         0.0.0.0
1                     192.168.1.1
2    2001:db8:85a3::8a2e:370:7334
dtype: ip

Note the dtype in that output. That’s a custom dtype (like category) defined outside pandas.

We register a custom accessor with pandas claiming the .ip namespace (just like pandas uses .str or .dt or .cat):

In [8]: ser.ip.isna
Out[8]:
0     True
1    False
2    False
dtype: bool

In [9]: ser.ip.is_ipv6
Out[9]:
0    False
1    False
2     True
dtype: bool

I’m extremely interested in seeing what the community builds on top of this interface. Joris has already tested out the Cythonized geopandas extension, which stores a NumPy array of pointers to geometry objects, and things seem great. I could see someone (perhaps you, dear reader?) building a JSONArray array type for working with nested data. That combined with custom .json accessor, perhaps with a jq-like query language should make for a powerful combination.

I’m also happy to have to say “Closed, out of scope; sorry.” less often. Now it can be “Closed, out of scope; do it outside of pandas.” :)

Open Source Success Story

It’s worth taking a moment to realize that this was a great example of open source at its best.

1. A company had a need for a tool. They didn’t have the expertise or desire to build and maintain it internally, so they approached Anaconda (a for-profit company with a great OSS tradition) to do it for them.
2. A proposal was made and rejected by the pandas community. You can’t just “buy” features in pandas if it conflicts too strongly with the long-term goals for the project.
3. A more general solution was found, with minimal changes to pandas itself, allowing anyone to do this type of extension outside of pandas.
4. We built the cyberpandas, which to users will feel like a first-class array type in pandas.

Thanks to the tireless reviews from the other pandas contributors, especially Jeff Reback, Joris van den Bossche, and Stephen Hoyer. Look forward to these changes in the next major pandas release.

This work is supported by Anaconda Inc and the Data Driven Discovery Initiative from the Moore Foundation.

This past week, I had a chance to visit some of the scikit-learn developers at Inria in Paris. It was a fun and productive week, and I’m thankful to them for hosting me and Anaconda for sending me there. This article will talk about some improvements we made to improve training scikit-learn models using a cluster.

Scikit-learn uses joblib for simple parallelism in many places. Anywhere you pass an n_jobs keyword, scikit-learn is internally calling joblib.Parallel(...), and doing a batch of work in parallel. The estimator may have an embarrassingly parallel step internally (fitting each of the trees in a RandomForest for example). Or your meta-estimator like GridSearchCV may try out many combinations of hyper-parameters in parallel.

You can think of joblib as a broker between the user and the algorithm author. The user comes along and says, “I have n_jobs cores, please use them!”. Scikit-Learn says “I have all these embarrassingly parallel tasks to be run as part of fitting this estimator.” Joblib accepts the cores from the user and the tasks from scikit-learn, runs the tasks on the cores, and hands the completed tasks back to scikit-learn.

Joblib offers a few “backends” for how to do your parallelism, but they all boil down to using many processes versus using many threads.

Parallelism in Python

A quick digression on single-machine parallelism in Python. We can’t say up front that using threads is always better or worse than using processes. Unfortunately the relative performance depends on the specific workload. But we do have some general heuristics that come down to serialization overhead and Python’s Global Interpreter Lock (GIL).

The GIL is part of CPython, the C program that interprets and runs your Python program. It limits your Python process so that only one thread is executing Python at once, defeating your parallelism. Fortunately, much of the numerical Python stack is written in C, Cython, C++, Fortran, or numba, and may be able to release the GIL. This means your “Python” program, which is calling into Cython or C via NumPy or pandas, can get real thread-based parallelism without being limited by the GIL. The main caveat here that manipulating strings or Python objects (lists, dicts, sets, etc) typically requires holding the GIL.

So, if we have the option of choosing threads or processes, which do we want? For most numeric / scientific workloads, threads are better than processes because of shared memory. Each thread in a thread-pool can view (and modify!) the same large NumPy array. With multiple processes, data must be serialized between processes (perhaps using pickle). For large arrays or dataframes this can be slow, and it may blow up your memory if the data a decent fraction of your machine’s RAM. You’ll have a full copy in each processes.

See Matthew Rocklin’s article and David Beazley’s page if you want to learn more.

Distributed Training with dask.distributed

For a while now, you’ve been able to use dask.distributed as a backend for joblib. This means that in most places scikit-learn offers an n_jobs keyword, you’re able to do the parallel computation on your cluster.

This is great when

1. Your dataset is not too large (since the data must be sent to each worker)
2. The runtime of each task is long enough that the overhead of serializing the data across the network to the worker doesn’t dominate the runtime
3. You have many parallel tasks to run (else, you’d just use a local thread or process pool and avoid the network delay)

Fitting a RandomForest is a good example of this. Each tree in a forest may be built independently of every other tree. This next code chunk shows how you can parallelize fitting a RandomForestClassifier across a cluster, though as discussed later this won’t work on the currently released versions of scikit-learn and joblib.

from sklearn.externals import joblib
from dask.distributed import Client
import distributed.joblib  # register the joblib backend

client = Client('dask-scheduler:8786')

with joblib.parallel_backend("dask", scatter=[X_train, y_train]):
    clf.fit(X_train, y_train)

The .fit call is parallelized across all the workers in your cluster. Here’s the distributed dashboard during that training.

The center pane shows the task stream as they complete. Each rectangle is a single task, building a single tree in a random forest in this case. Workers are represented vertically. My cluster had 8 workers with 4 cores each, which means up to 32 tasks can be processed simultaneously. We fit the 200 trees in about 20 seconds.

Changes to Joblib

Above, I said that distributed training worked in most places in scikit-learn. Getting it to work everywhere required a bit more work, and was part of last week’s focus.

First, dask.distributed’s joblib backend didn’t handle nested parallelism well. This may occur if you do something like

gs = GridSearchCV(Estimator(n_jobs=-1), n_jobs=-1)
gs.fit(X, y)

Previously, that caused deadlocks. Inside GridSearchCV, there’s a call like

# In GridSearchCV.fit, the outer layer
results = joblib.Parallel(n_jobs=n_jobs)(fit_estimator)(...)

where fit_estimator is a function that itself tries to do things in parallel

# In fit_estimator, the inner layer
results = joblib.Parallel(n_jobs=n_jobs)(fit_one)(...)

So the outer level kicks off a bunch of joblib.Parallel calls, and waits around for the results. For each of those Parallel calls, the inner level tries to make a bunch of joblib.Parallel calls. When joblib tried to start the inner ones, it would ask the distributed scheduler for a free worker. But all the workers were “busy” waiting around for the outer Parallel calls to finish, which weren’t progressing because there weren’t any free workers! Deadlock!

dask.distributed has a solution for this case (workers secede from the thread pool when they start a long-running Parllel call, and rejoin when they’re done), but we needed a way to negotiate with joblib about when the secede and rejoin should happen. Joblib now has an API for backends to control some setup and teardown around the actual function execution. This work was done in Joblib #538 and dask-distributed #1705.

Second, some places in scikit-learn hard-code the backend they want to use in their Parallel() call, meaning the cluster isn’t used. This may be because the algorithm author knows that one backend performs better than others. For example, RandomForest.fit performs better with threads, since it’s purely numeric and releases the GIL. In this case we would say the Parallel call prefers threads, since you’d get the same result with processes, it’d just be slower.

Another reason for hard-coding the backend is if the correctness of the implementation relies on it. For example, RandomForest.predict preallocates the output array and mutates it from many threads (it knows not to mutate the same place from multiple threads). In this case, we’d say the Parallel call requires shared memory, because you’d get an incorrect result using processes.

The solution was to enhance joblib.Parallel to take two new keywords, prefer and require. If a Parallel call prefers threads, it’ll use them, unless it’s in a context saying “use this backend instead”, like

def fit(n_jobs=-1):
    return joblib.Parallel(n_jobs=n_jobs, prefer="threads")(...)


with joblib.parallel_backend('dask'):
    # This uses dask's workers, not threads
    fit()

On the other hand, if a Parallel requires a specific backend, it’ll get it.

def fit(n_jobs=-1):
    return joblib.Parallel(n_jobs=n_jobs, require="sharedmem")(...)

with joblib.parallel_backend('dask'):
    # This uses the threading backend, since shared memory is required
    fit()

This is a elegant way to negotiate a compromise between

1. The user, who knows best about what resources are available, as specified by the joblib.parallel_backend context manager. And,
2. The algorithm author, who knows best about the GIL handling and shared memory requirements.

This work was done in Joblib #602.

After the next joblib release, scikit-learn will be updated to use these options in places where the backend is currently hard-coded. My example above used a branch with those changes.

Look forward for these changes in the upcoming joblib, dask, and scikit-learn releases. As always, let me know if you have any feedback.

This past week, I had a chance to visit some of the scikit-learn developers at Inria in Paris. It was a fun and productive week, and I’m thankful to them for hosting me and Anaconda for sending me there.

Towards the end of our week, Gael threw out the observation that for many applications, you don’t need to train on the entire dataset, a sample is often sufficient. But it’d be nice if the trained estimator would be able to transform and predict for dask arrays, getting all the nice distributed parallelism and memory management dask brings.

This intrigued me, and I had a 9 hour plane ride, so…

dask_ml.iid

I put together the dask_ml.iid sub-package. The estimators contained within are appropriate for data that are independently and identically distributed (IID). Roughly speaking, your data is IID if there aren’t any “patterns” in the data as you move top to bottom. For example, time-series data is often not IID, there’s often an underlying time trend to the data. Or the data may be autocorrelated (if y was above average yesterday, it’ll probably be above average today too). If your data is sorted, say by customer ID, then it likely isn’t IID. You might be able to shuffle it in this case.

If your data are IID, it may be OK to just fit the model on the first block. In principal, it should be a representative sample of your entire dataset.

Here’s a quick example. We’ll fit a GradientBoostingClassifier. The dataset will be 1,000,000 x 20, in chunks of 10,000. This would take way too long to fit regularly. But, with IID data, we may be OK fitting the model on just the the first 10,000 observations.

>>> from dask_ml.datasets import make_classification
>>> from dask_ml.iid.ensemble import GradientBoostingClassifier

>>> X, y = make_classification(n_samples=1_000_000, chunks=10_100)

>>> clf = GradientBoostingClassifier()
>>> clf.fit(X, y)

At this point, we have a scikit-learn estimator that can be used to transform or predict for dask arrays (in parallel, out of core or distributed across your cluster).

>>> prob_a
dask.array<predict_proba, shape=(1000000, 2), dtype=float64, chunksize=(10000, 2)>

>>> prob_a[:10].compute()
array([[0.98268198, 0.01731802],
       [0.41509521, 0.58490479],
       [0.97702961, 0.02297039],
       [0.91652623, 0.08347377],
       [0.96530773, 0.03469227],
       [0.94015097, 0.05984903],
       [0.98167384, 0.01832616],
       [0.97621963, 0.02378037],
       [0.95951444, 0.04048556],
       [0.98654415, 0.01345585]])

An alternative to dask_ml.iid is to sample your data and use a regular scikit-learn estimator. But the dask_ml.iid approach is slightly preferable, since post-fit tasks like prediction can be done on dask arrays in parallel (and potentially distributed). Scikit-Learn’s estimators are not dask-aware, so they’d just convert it to a NumPy array, possibly blowing up your memory.

If dask and dask_ml.iid had existed a few years ago, it would have solved all the “big data” needs of my old job. Personally, I never hit a problem where, if my dataset was already large, training on an even larger dataset was the answer. I’d always hit the level part of the learning curve, or was already dealing with highly imbalanced classes. But, I would often have to make predictions for a much larger dataset. For example, I might have trained a model on “all the customers for this store” and predicted for “All the people in Iowa”.

Today we released the first version of dask-ml, a library for parallel and distributed machine learning. Read the documentation or install it with

pip install dask-ml

Packages are currently building for conda-forge, and will be up later today.

conda install -c conda-forge dask-ml

The Goals

dask is, to quote the docs, “a flexible parallel computing library for analytic computing.” dask.array and dask.dataframe have done a great job scaling NumPy arrays and pandas dataframes; dask-ml hopes to do the same in the machine learning domain.

Put simply, we want

est = MyEstimator()
est.fit(X, y)

to work well in parallel and potentially distributed across a cluster. dask provides us with the building blocks to do that.

What’s Been Done

dask-ml collects some efforts that others already built:

• distributed joblib: scaling out some scikit-learn operations to clusters (from distributed.joblib)
• hyper-parameter search: Some drop in replacements for scikit-learn’s GridSearchCV and RandomizedSearchCV classes (from dask-searchcv)
• distributed GLMs: Fit large Generalized Linear Models on your cluster (from dask-glm)
• dask + xgboost: Peer a dask.distributed cluster with XGBoost running in distributed mode (from dask-xgboost)
• dask + tensorflow: Peer a dask.distributed cluster with TensorFlow running in distributed mode (from dask-tensorflow)
• Out-of-core learning in pipelines: Reuse scikit-learn’s out-of-core .partial_fit API in pipelines (from dask.array.learn)

In addition to providing a single home for these existing efforts, we’ve implemented some algorithms that are designed to run in parallel and distributed across a cluster.

• KMeans: Uses the k-means|| algorithm for initialization, and a parallelized Lloyd’s algorithm for the EM step.
• Preprocessing: These are estimators that can be dropped into scikit-learn Pipelines, but they operate in parallel on dask collections. They’ll work well on datasets in distributed memory, and may be faster for NumPy arrays (depending on the overhead from parallelizing, and whether or not the scikit-learn implementation is already parallel).

Help Contribute!

Scikit-learn is a robust, mature, stable library. dask-ml is… not. Which means there are plenty of places to contribute! Dask makes writing parallel and distributed implementations of algorithms fun. For the most part, you don’t even have to think about “where’s my data? How do I parallelize this?” Dask does that for you.

Have a look at the issues or propose a new one. I’d love to hear issues that you’ve run into when scaling the “traditional” scientific python stack out to larger problems.

This work is supported by Anaconda, Inc. and the Data Driven Discovery Initiative from the Moore Foundation.

This is part three of my series on scalable machine learning.

• Small Fit, Big Predict
• Scikit-Learn Partial Fit
• Parallel Machine Learning

You can download a notebook of this post [here][notebook].

In part one, I talked about the type of constraints that push us to parallelize or distribute a machine learning workload. Today, we’ll be talking about the second constraint, “I’m constrained by time, and would like to fit more models at once, by using all the cores of my laptop, or all the machines in my cluster”.

An Aside on Parallelism

In the case of Python, we have two main avenues of parallelization (which we’ll roughly define as using multiple “workers” to do some “work” in less time). Those two avenues are

1. multi-threading
2. multi-processing

For python, the most important differences are that

1. multi-threaded code can potentially be limited by the GIL
2. multi-processing code requires that data be serialized between processes

The GIL is the “Global Interpreter Lock”, an implementation detail of CPython that means only one thread in your python process can be executing python code at once.

This talk by Python core-developer Raymond Hettinger does a good job summarizing things for Python, with an important caveat: much of what he says about the GIL doesn’t apply to the scientific python stack. NumPy, scikit-learn, and much of pandas release the GIL and can run multi-threaded, using shared memory and so avoiding serialization costs. I’ll highlight his quote, which summarizes the situation:

Your weakness is your strength, and your strength is your weakness

The strength of threads is shared state. The weakness of threads is shared state.

Another wrinkle here is that when you move to a distributed cluster, you have to have multiple processes. And communication between processes becomes even more expensive since you’ll have network overhead to worry about, in addition to the serialization costs.

Fortunately, modules like concurrent.futures and libraries like dask make it easy to swap one mode in for another. Let’s make a little dask array:

import dask.array as da
import dask
import dask.threaded
import dask.multiprocessing

X = da.random.uniform(size=(10000, 10), chunks=(1000, 10))
result = X / (X.T @ X).sum(1)

We can swap out the scheduler with a context-manager:

%%time
with dask.set_options(get=dask.threaded.get):
    # threaded is the default for dask.array anyway
    result.compute()

%%time
with dask.set_options(get=dask.multiprocessing.get):
    result.compute()

Every dask collection (dask.array, dask.dataframe, dask.bag) has a default scheduler that typically works well for the kinds of operations it does. For dask.array and dask.dataframe, the shared-memory threaded scheduler is used.

Cost Models

In this talk, Simon Peyton Jones talks about parallel and distributed computing for Haskell. He stressed repeatedly that there’s no silver bullet when it comes to parallelism. The type of parallelism appropriate for a web server, say, may be different than the type of parallelism appropriate for a machine learning algorithm.

I mention all this, since we’re about to talk about parallel machine learning. In general, for small data and many models you’ll want to use the threaded scheduler. For bigger data (larger than memory), you’ll want want to use the distributed scheduler. Assuming the underlying NumPy, SciPy, scikit-learn, or pandas operation releases the GIL, you’ll be able to get nice speedups without the cost of serialization. But again, there isn’t a silver bullet here, and the best type of parallelism will depend on your particular problem.

Where to Parallelize

In a typical machine-learning workflow, there are typically ample opportunities for parallelism.

1. Over Hyper-parameters (one fit per combination of parameters)
2. Over Cross-validation folds (one fit per fold)
3. Within an algorithm (for some algorithms)

Scikit-learn already uses parallelism in many places, anywhere you see an n_jobs keyword.

This work is supported by Anaconda, Inc. and the Data Driven Discovery Initiative from the Moore Foundation.

This is part two of my series on scalable machine learning.

• Small Fit, Big Predict
• Scikit-Learn Partial Fit

You can download a notebook of this post here.

Scikit-learn supports out-of-core learning (fitting a model on a dataset that doesn’t fit in RAM), through it’s partial_fit API. See here.

The basic idea is that, for certain estimators, learning can be done in batches. The estimator will see a batch, and then incrementally update whatever it’s learning (the coefficients, for example). This link has a list of the algorithms that implement partial_fit.

Unfortunately, the partial_fit API doesn’t play that nicely with my favorite part of scikit-learn, pipelines, which we discussed at length in part 1. For pipelines to work, you would essentially need every step in the pipeline to have an out-of-core partial_fit version, which isn’t really feasible; some algorithms just have to see the entire dataset at once. Setting that aside, it wouldn’t be great for a user, since working with generators of datasets is awkward compared to the expressivity we get from pandas and NumPy.

Fortunately, we have great data containers for larger than memory arrays and dataframes: dask.array and dask.dataframe. We can

1. Use dask for pre-processing data in an out-of-core manner
2. Use scikit-learn to fit the actual model, out-of-core, using the partial_fit API

And with a little bit of work, all of this can be done in a pipeline. The rest of this post shows how.

Big Arrays

If you follow along in the companion notebook, you’ll see that I generate a dataset, replicate it 100 times, and write the results out to disk. I then read it back in as a pair of dask.dataframes and convert them to a pair of dask.arrays. I’ll skip those details to focus on main goal: using sklearn.Pipelines on larger-than-memory datasets. Suffice to say, we have a function read that gives us our big X and y:

X, y = read()
X

dask.array<concatenate, shape=(100000000, 20), dtype=float64, chunksize=(500000, 20)>

y

dask.array<squeeze, shape=(100000000,), dtype=float64, chunksize=(500000,)>

So X is a 100,000,000 x 20 array of floats (I have float64s, you’re probably fine with float32s) that we’ll use to predict y. I generated the dataset, so I know that y is either 0 or 1. We’ll be doing classification.

(X.nbytes + y.nbytes) / 10**9

16.8

My laptop has 16 GB of RAM, and the dataset is 16.8 GB. We can’t simply read the entire thing into memory. We’ll use dask for the preprocessing, and scikit-learn for the fitting. We’ll have a small pipeline

1. Scale the features by mean and variance
2. Fit an SGDClassifier

I’ve implemented a daskml.preprocessing.StandardScaler, using dask, in about 40 lines of code (see here). The scaling will be done completely in parallel and completely out-of-core.

I haven’t implemented a custom SGDClassifier, because that’d be much more than 40 lines of code. Instead, I’ve put together a small wrapper that will use scikit-learn’s SGDClassifier.partial_fit to fit the model out-of-core (but not in parallel).

from daskml.preprocessing import StandardScaler
from daskml.linear_model import BigSGDClassifier  # The wrapper

from dask.diagnostics import ResourceProfiler, Profiler, ProgressBar
from sklearn.pipeline import make_pipeline

As a user, the API is the same as scikit-learn. Indeed, it is just a regular sklearn.pipeline.Pipeline.

pipe = make_pipeline(
    StandardScaler(),
    BigSGDClassifier(classes=[0, 1], max_iter=1000, tol=1e-3, random_state=2),
)

And fitting is identical as well: pipe.fit(X, y). We’ll collect some performance metrics as well, so we can analyze our parallelism.

%%time
rp = ResourceProfiler()
p = Profiler()


with p, rp:
    pipe.fit(X, y)

CPU times: user 2min 38s, sys: 1min 44s, total: 4min 22s
Wall time: 1min 47s

And that’s it. It’s just a regular scikit-learn pipeline, operating on a larger-than-memory data. pipe has has all the regular methods you would expect, predict, predict_proba, etc. You can get to the individual attributes like pipe.steps[1][1].coef_.

One important point to stress here: when we get to the BigSGDClassifier.fit at the end of the pipeline, everything is done serially. We can see that by plotting the Profiler we captured up above:

That graph shows the tasks (the rectangles) each worker (a core on my laptop) executed over time. Workers are along the vertical axis, and time is along the horizontal. Towards the start, when we’re reading off disk, converting to dask.arrays, and doing the StandardScaler, everything is in parallel. Once we get to the BigSGDClassifier, which is just a simple wrapper around sklearn.linear_model.SGDClassifier, we lose all our parallelism*.

The predict step is done entirely in parallel.

with rp, p:
    predictions = pipe.predict(X)
    predictions.to_dask_dataframe(columns='a').to_parquet('predictions.parq')

That took about 40 seconds, from disk to prediction, and back to disk on 16 GB of data, using all 8 cores of my laptop.

How?

When I had this idea last week, of feeding blocks of dask.array to a scikit-learn estimator’s partial_fit method, I thought it was pretty neat. Turns out Matt Rocklin already had the idea, and implemented it in dask, two years ago.

Roughly speaking, the implementation is:

class BigSGDClassifier(SGDClassifier):
    ...
    
    def fit(self, X, y):
        # ... some setup
        for xx, yy in by_blocks(X, y):
            self.partial_fit(xx, yy)
        return self

If you aren’t familiar with dask, its arrays are composed of many smaller NumPy arrays (blocks in the larger dask array). We iterate over the dask arrays block-wise, and pass them into the estimators partial_fit method. That’s exactly what you would be doing if you were using, say, a generator feed NumPy arrays to the partial_fit method. Only you can manipulate a dask.array like regular NumPy array, so things are more convenient.

Some Challenges

For our small pipeline, we had to make two passes over the data. One to fit the StandardScaler and one to fit the BigSGDClassifier. In general, with this approach, we’ll have to make one pass per stage of the pipeline, which isn’t great. I think this is unavoidable with the current design, but I’m considering ways around it.

Recap

We’ve seen a way to use scikit-learn’s existing estimators on larger-than-memory dask arrays by passing the blocks of a dask array to the partial_fit method. This enables us to use Pipelines on larger-than-memory datasets.

Let me know what you think. I’m pretty excited about this because it removes some of the friction around using sckit-learn Pipelines with out-of-core estimators. In dask-ml, I’ve implemented similar wrappers for

• SGDRegressor
• PassiveAggressiveClassifier
• PassiveAggressiveRegressor
• Perceptron
• MPLClassifier
• MLPRegressor
• MiniBatchKMeans

I’ll be packaging this up in daskml to make it more usable for the community over the next couple weeks. If this type of work interests you, please reach out on Twitter or by email at mailto:[email protected]. If you’re interested in contributing, I think a library of basic transformers that operate on NumPy and dask arrays and pandas and dask DataFrames would be extremely useful. I’ve started an issue to track this progress. Contributions would be extremely welcome.

Next time we’ll be going back to smaller datasets. We’ll see how dask can help us parallelize our work to fit more models in less time.

This work is supported by Anaconda Inc. and the Data Driven Discovery Initiative from the Moore Foundation.

Anaconda is interested in scaling the scientific python ecosystem. My current focus is on out-of-core, parallel, and distributed machine learning. This series of posts will introduce those concepts, explore what we have available today, and track the community’s efforts to push the boundaries.

You can download a Jupyter notebook demonstrating the analysis here.

Constraints

I am (or was, anyway) an economist, and economists like to think in terms of constraints. How are we constrained by scale? The two main ones I can think of are

1. I’m constrained by size: My training dataset fits in RAM, but I have to predict for a much larger dataset. Or, my training dataset doesn’t even fit in RAM. I’d like to scale out by adopting algorithms that work in batches locally, or on a distributed cluster.
2. I’m constrained by time: I’d like to fit more models (think hyper-parameter optimization or ensemble learning) on my dataset in a given amount of time. I’d like to scale out by fitting more models in parallel, either on my laptop by using more cores, or on a cluster.

These aren’t mutually exclusive or exhaustive, but they should serve as a nice framework for our discussion. I’ll be showing where the usual pandas + scikit-learn for in-memory analytics workflow breaks down, and offer some solutions for scaling out to larger problems.

This post will focus on cases where your training dataset fits in memory, but you must predict on a dataset that’s larger than memory. Later posts will explore into parallel, out-of-core, and distributed training of machine learning models.

Don’t forget your Statistics

Statistics is a thing¹. Statisticians have thought a lot about things like sampling and the variance of estimators. So it’s worth stating up front that you may be able to just

SELECT *
FROM dataset
ORDER BY random()
LIMIT 10000;

and fit your model on a (representative) subset of your data. You may not need distributed machine learning. The tricky thing is selecting how large your sample should be. The “correct” value depends on the complexity of your learning task, the complexity of your model, and the nature of your data. The best you can do here is think carefully about your problem and to plot the learning curve.

As usual, the scikit-learn developers do a great job explaining the concept in addition to providing a great library. I encourage you to follow that link. This gist is that—for some models on some datasets—training the model on more observations doesn’t improve performance. At some point the learning curve levels off and you’re just wasting time and money training on those extra observations.

For today, we’ll assume that we’re on the flat part of the learning curve. Later in the series we’ll explore cases where we run out of RAM before the learning curve levels off.

Fit, Predict

In my experience, the first place I bump into RAM constraints is when my training dataset fits in memory, but I have to make predictions for a dataset that’s orders of magnitude larger. In these cases, I fit my model like normal, and do my predictions out-of-core (without reading the full dataset into memory at once).

We’ll see that the training side is completely normal (since everything fits in RAM). We’ll see that dask let’s us write normal-looking pandas and NumPy code, so we don’t have to worry about writing the batching code ourself.

To make this concrete, we’ll use the (tried and true) New York City taxi dataset. The goal will be to predict if the passenger leaves a tip. We’ll train the model on a single month’s worth of data (which fits in my laptop’s RAM), and predict on the full dataset².

Let’s load in the first month of data from disk:

dtype = {
    'vendor_name': 'category',
    'Payment_Type': 'category',
}

df = pd.read_csv("data/yellow_tripdata_2009-01.csv", dtype=dtype,
                 parse_dates=['Trip_Pickup_DateTime', 'Trip_Dropoff_DateTime'],)
df.head()

The January 2009 file has about 14M rows, and pandas takes about a minute to read the CSV into memory. We’ll do the usual train-test split:

X = df.drop("Tip_Amt", axis=1)
y = df['Tip_Amt'] > 0

X_train, X_test, y_train, y_test = train_test_split(X, y)

print("Train:", len(X_train))
print("Test: ", len(X_test))

Train: 10569309
Test:  3523104

Aside on Pipelines

The first time you’re introduced to scikit-learn, you’ll typically be shown how you pass two NumPy arrays X and y straight into an estimator’s .fit method.

from sklearn.linear_model import LinearRegression

est = LinearRegression()
est.fit(X, y)

Eventually, you might want to use some of scikit-learn’s pre-processing methods. For example, we might impute missing values with the median and normalize the data before handing it off to LinearRegression. You could do this “by hand”:

from sklearn.preprocessing import Imputer, StandardScaler

imputer = Imputer(strategy='median')
X_filled = imputer.fit_transform(X, y)

scaler = StandardScaler()
X_scaled = X_scaler.fit_transform(X_filled, y)

est = LinearRegression()
est.fit(X_scaled, y)

We set up each step, and manually pass the data through: X -> X_filled -> X_scaled.

The downside of this approach is that we now have to remember which pre-processing steps we did, and in what order. The pipeline from raw data to fit model is spread across multiple python objects. A better approach is to use scikit-learn’s Pipeline object.

from sklearn.pipeline import make_pipeline

pipe = make_pipeline(
    Imputer(strategy='median'),
    StandardScaler(),
    LinearRegression()
)
pipe.fit(X, y)

Each step in the pipeline implements the fit, transform, and fit_transform methods. Scikit-learn takes care of shepherding the data through the various transforms, and finally to the estimator at the end. Pipelines have many benefits but the main one for our purpose today is that it packages our entire task into a single python object. Later on, our predict step will be a single function call, which makes scaling out to the entire dataset extremely convenient.

If you want more information on Pipelines, check out the scikit-learn docs, this blog post, and my talk from PyData Chicago 2016. We’ll be implementing some custom ones, which is not the point of this post. Don’t get lost in the weeds here, I only include this section for completeness.

Our Pipeline

This isn’t a perfectly clean dataset, which is nice because it gives us a chance to demonstrate some pandas’ pre-processing prowess, before we hand the data of to scikit-learn to fit the model.

from sklearn.pipeline import make_pipeline
# We'll use FunctionTransformer for simple transforms
from sklearn.preprocessing import FunctionTransformer
# TransformerMixin gives us fit_transform for free
from sklearn.base import TransformerMixin

There are some minor differences in the spelling on “Payment Type”:

df.Payment_Type.cat.categories

Index(['CASH', 'CREDIT', 'Cash', 'Credit', 'Dispute', 'No Charge'], dtype='object')

We’ll reconcile that by lower-casing everything with a .str.lower(). But resist the temptation to just do that imperatively inplace! We’ll package it up into a function that will later be wrapped up in a FunctionTransformer.

def payment_lowerer(X):
    return X.assign(Payment_Type=X.Payment_Type.str.lower())

I should note here that I’m using .assign to update the variables since it implicitly copies the data. We don’t want to be modifying the caller’s data without their consent.

Not all the columns look useful. We could have easily solved this by only reading in the data that we’re actually going to use, but let’s solve it now with another simple transformer:

class ColumnSelector(TransformerMixin):
    "Select `columns` from `X`"
    def __init__(self, columns):
        self.columns = columns

    def fit(self, X, y=None):
        return self

    def transform(self, X, y=None):
        return X[self.columns]

Internally, pandas stores datetimes like Trip_Pickup_DateTime as a 64-bit integer representing the nanoseconds since some time in the 1600s. If we left this untransformed, scikit-learn would happily transform that column to its integer representation, which may not be the most meaningful item to stick in a linear model for predicting tips. A better feature might the hour of the day:

class HourExtractor(TransformerMixin):
    "Transform each datetime in `columns` to integer hour of the day"
    def __init__(self, columns):
        self.columns = columns

    def fit(self, X, y=None):
        return self

    def transform(self, X, y=None):
        return X.assign(**{col: lambda x: x[col].dt.hour
                           for col in self.columns})

Likewise, we’ll need to ensure the categorical variables (in a statistical sense) are categorical dtype (in a pandas sense). We want categoricals so that we can call get_dummies later on without worrying about missing or extra categories in a subset of the data throwing off our linear algebra (See my talk for more details).

class CategoricalEncoder(TransformerMixin):
    """
    Convert to Categorical with specific `categories`.
    
    Examples
    --------
    >>> CategoricalEncoder({"A": ['a', 'b', 'c']}).fit_transform(
    ...     pd.DataFrame({"A": ['a', 'b', 'a', 'a']})
    ... )['A']
    0    a
    1    b
    2    a
    3    a
    Name: A, dtype: category
    Categories (2, object): [a, b, c]
    """
    def __init__(self, categories):
        self.categories = categories
        
    def fit(self, X, y=None):
        return self
        
    def transform(self, X, y=None):
        X = X.copy()
        for col, categories in self.categories.items():
            X[col] = X[col].astype('category').cat.set_categories(categories)
        return X