Pragmatic ways of using Rust in your data project

One common strategy is to wrap the Rust part as a Python extension module. With enough care, the extensions module can have a pythonic feel and substantially improve performance. While libraries, such as PyO3, offer streamlined APIs, this task can still require lot of work.

An often simpler alternative is to package the Rust part as an executable and communicate via files or network. This talk will focus on JSON messages exchanged via stdin / stdout or dataframe-like data in Arrow-compatible files. JSON is broadly supported in both Python and Rust and serialization can easily be handled with libraries such as SerDe (Rust) or cattrs (Python). The Arrow in-memory format supports complex data types, such as structs, lists, maps, or unions. These files can then be efficiently processed in Python by ever an growing list of libraries, most prominently pandas and polars.

I will discuss the different strategies using real-world use cases and offer tips on how to implement them. Finally I will end by summarizing the respective strengths and weaknesses of the approaches.

This session took place in track Data Handling and was classified suitable for novice domain / intermediate python by the speaker.

Transcript (auto)

Auto-generated from the recording utilizing Open-Source AI. Speaker labels (Speaker 1, Speaker 2) reflect diarization, not identity. Timestamps refer to the recording.

Speaker 1 [00:03]

Cool. Thank you very much for the introduction. Thank you all for being here. And as the title states, I would like to talk a little bit about Rust with you, but doing it in a pragmatic way. So don't expect too much high-level tech or cool tech. This will be really down-to-earth, hopefully. And maybe a little bit about me. I'm a physicist by training and then essentially started working as a data scientist. But now at Volkswagen, and I've been using Python for actually quite some time now. and maybe to give you also a little bit of a background why I really like Python is essentially this interactivity right you can see the results of what you're doing when you're building a model when you're analyzing data when you but also when you're writing tests for example something I like to do and then you can really see the outcome and then iterate on that however there's also a little bit of a dark side to Python I would say specifically Python can be very very slow if If your data doesn't fit into NumPy, Pandas, it can be slow. And I think much more importantly also, it's really easy to leave the golden path. Like you make one wrong step and suddenly it becomes slow. And I think this is something that is a little bit tricky with Python. And I would like to present you some strategies of maybe using Rust to get around these rough edges. And the strategy will not be to replace Python. I really like Python. Python is a cool language. And for everything that is interactive, I think it's also unbeatable. at least for me but rather I would like to show you a little bit how to sprinkle in Rust here and there and then how to make your code fast overall and actually I would say that Python and Rust have complementary strengths where one is really interactive really nice to like play around with the other is rather slow and also in compile times as we saw this morning but then it's fast and the other is a bit slower if you do maybe non-standard stuff. so i'm a data scientist so what i typically do is i start somehow with some kind of data let's call it complex data think like a giant heap of json or pdf documents as we will see in a second and then i typically want to put it in the data frame as i as we do as you do i guess and then perform some kind of analysis on top of it build a model and maybe already here like a like an advanced notice, so to say, what you can by now put in the data frame also is really much more complex than the usual. You have a string column, you have an integer column. But with the arrow format that we already heard a couple of times in this conference, you can do really cool things. In principle, whatever you can think of, you can put by now in a data frame. So really, data frame is more than just a simple table. And specifically, I would like to talk today about two examples. One is parsing PDFs. Here I chose bank statements because I have them available on my laptop. And then also how to convert a type of JSON files into a data frame than to do analysis on top of. And here this is a JSON data set by Spotify, the million playlist data set. But we will get to all the details in a bit. So and now I would like to show you a little bit how you can use Rust in these two tasks and maybe some strategies. Like a quick warning, there will be Rust code ahead. I don't expect you to be familiar with Rust. So maybe you don't get every single line. My goal here is more to give you some pointers and hints and give you a feeling how it looks like, not that you understand, like I said, every single line. So don't worry. Hopefully it's kind of self-explanatory. And if not, like sleep for 30 seconds and then some Python code will come again. So what do I like about Rust? So by the clear, the motivation here is like for performance and memory efficiency. I want something that is fast and efficient, and Rust delivers on that. I think that's also one of the main goals of when Rust was designed. The other thing is that I particularly like about Rust is that it's very easy to integrate into other things. It doesn't have a runtime. It doesn't have a garbage collector. it's really built to interface with anything that speaks CAVI, in particular Python as well. So you can use the same Rust library in Java, in Python, in Go, and you don't, it's really easy. Whereas with like something like Go, you have a garbage collector, suddenly it becomes more complicated. Also, I really like how Rust is designed. It has some kind of features that fit nicely together and form a cohesive whole that you can use to build really high-level level interfaces, and I will show you a couple of examples later on. And finally, as we also heard this morning, Rust comes with great tooling, great documentation. Not every library is well documented, but overall I would say the quality is very high, in particular of the core project. So this morning we heard about how to use PIO3 to build essentially something like some Rust library that you can use in Python. And I would definitely say this is the gold standard. If you want to build a high-quality package that you maybe want to deliver to other people, you should definitely check it out. I think that's the way to go. And maybe sometimes if you check out some open source package that you're using and then you're interested how it's built, it may come to some surprise to you that if you look at the source code, there's not a single line of Python or not a single Python file. And PyO3 really makes it easy to build Pythonic interfaces on top of your Rust code without writing Python. And I think this is really like a testament to the library, to the quality. And this is an example, a very simple example, it's just from the documentation, essentially a function that takes in a number, an integer, and then doubles it. And one thing that is very specific about Rust, and that will come up a couple of times later on, is that Rust allows you to do code generation via those macros. So every time you see this hash and a bracket, or maybe something with an exclamation mark, This is code generation where the compiler generates code for you that you don't have to write yourself. And this also makes Rust very nice to use and very short. Like I said, we will not or maybe I haven't said this already, but this will not be about PIO3 actually. I would rather want to advertise for some hacky easy way, specifically let's build our own command line tools and then interface with these command line tools to get essentially the speed up. And as an example, let's stick with the whole doubling number thing. So we have this command line tool called IOPatterns-double.exe. So the name is from the repository that is also on GitHub if you're interested. And then the command line tools takes the JSON in with a value and then returns another JSON out with a value doubled. And why JSON? It's the lingua franca of data exchange, I would say. Everything speaks JSON, everything writes JSON, so it's nice and easy. If you want to have a bit more performance, maybe go for binary encoding. And there are also options then to generate the code that you need to essentially go from A to B. When you want to use this tool from Python or this command line tool, you can use the built-in subprocess module. So what you do is you call the command line tool, and then you need to configure a couple of options. Importantly, in Rust, every string is UTF-8 encoded, so you should set the encoding, whatever the encoding of your machine is and then we want to get the output out so you say capture output we want to have an exception if an error occurs check equals true and then we essentially we have the last two lines where we feed the input in and the output out and i would like to show you a little bit how you can easily then write the code for that so you get essentially this encoding of the input and the decoding of the output and the basic strategy is to use types and code generation both in Rust and in Python and for example here I used data classes depending on your choice maybe you would on what you like you probably use something different for me typically data classes works quite well from the standard library and so here we define essentially our interface as types and then there's a very nice library called C address that gives you for the simple cases a zero configuration way of going from these data classes to something that was JSON compatible and back again. If you are using unions, so something that could be a float or an int, then you need to write a bit more code yourself. You will find some examples in the repository of different coding schemes, but let's skip with a simple case for now. So if you go from the way from Python objects to JSON compatible objects, so to say, you can use the unstructure call. So what it does, it takes your data class in, and then essentially returns something that is like plain old Python dictionaries and lists. And then you can use JSON.dumps to get essentially something that is a JSON string. If you want to go the other way around, you just use structure, and essentially you load your JSON string, then you structure it into the output, so you get back your output data class that you then use in your code, so into domain objects. And that's essentially everything you need for the Python side. And all we did is essentially we needed to write the type definitions, the rest is taken care of by the type hints. In Rust actually we do the same. So here we define our types as well. In this case it looks a bit different because it's a different language. We have our structure, so this is essentially the data class parts. And here we are again using code generation, so everything with a hash. And we're using a library called Serder for serialization, deserialization, which is a de facto standard in the Rust ecosystem to essentially get types serialized and deserialized. And again, all the code that takes care of converting from JSON to your types and back again is generated for you. And then if you want to implement this IDOO doubling interface, what you can do then is essentially you read from standard in and then build your Rust object out. Then you perform the operation, in this case doubling the number, and then write it back out to send it out. That's it. And maybe to contrast let's say the two ways like PIO3 or some other bindings generator and doing it this custom CLI tool way. If you go PIO3 this is definitely the way to go if you want to build high quality extension modules also maybe ship something put it on PIPI I would probably go this way. Also if you go for maximum efficiency so if you want to ring out every last bit of performance, this is the way to go. However, there's also some drawbacks. I think you need to quite well understand how PIO3 works. I think you also need to understand the C API, because sometimes there are these rough cases where you actually need to understand how this is then translated into the C API, in particular if bugs happen. There's also this, I would say, non-trivial building and installation step. We've seen this morning that you can just use Maturin, but still you have this extra step where you need to build something, compile something, and then integrate it into your Python environment. One thing that is a big drawback for myself, personally speaking, you don't have reloading. So Python extension modules or native extension modules, as far as I'm aware, cannot be reloaded. So for example, if you use it in your Jupyter interpreter, you have to kill the interpreter and start from scratch, which is essentially something that I really don't like. When you go the CLI way about things, I would say it's really easy to get started. You have like JSON in, JSON out, you can inspect the results and if something goes wrong it should be easy to figure out. It's easy to build and distribute. It's much easier also to cross-compile, even though Rust is already great in cross-compiling, it gets even easier if you build executables, not shared libraries. However, it's not super efficient, so you clearly have performance overhead, and also So you don't have shared memory or shared objects. So it's more for functional patterns, function input in, output out. But you can always do refactor into PIO3 if you want to go for more performance. And after this little bit of explanation, I would actually show you how you can use it and that even though it kind of seems a little bit maybe too simplistic or too simplistic, it actually can get out some nice speed benefits from it. And the application specifically is about parsing PDF files. And I would like to tell you a little bit about why PDF files are actually a little bit tricky for Python. So if you have this fake PDF files, I didn't want to show you my bank statements, what you see is essentially the bank statement of some German bank. And what you then see in the PDF actually is instructions to some abstract machine that says something like move the text position to position one two and then at this position put this text and this can then be combined with like transformation matrices and maybe with color changes and so on and so on and what you now need to be do is when you want to like get the information out of your PDF you actually need an interpreter for the stream of commands and this is here taken from the PDFminer 6 library very nice library but now Now you have this interpreter running in an interpreter and this then gets slow. So this is where you can get the speed up. So what we will do is essentially we will go from the raw PDF to tuples or like some objects where we have position and text fragments and then we will use the second step to find the header and the footer of the transaction tables, group the different transaction items and then pass essentially the date and the amounts. For the first part, we will investigate Python and Rust. Whereas for the second part, this will always stay in Python. So essentially you have a very small part of your Python application that you replace and the large part will stay in Python SS. Actually as we will see, the Rust part is quite long. When I do that for one year of my bank accounts, we have essentially these two steps. Step one is something like 17 lines of PDF-6, mostly getting the type definitions and converting it to the types, and then you have something like 106 lines of Python to interpret the results. And for one year, this took on my laptop something like 25 seconds. So you have like 12 bank statements, and yeah, 25 seconds, I would say that's not great performance-wise. For the Rust side of things, we replaced this first step with a command line tool and then some interfacing tool. There we use a library called PDF extract, and what this looks like roughly is you have this command line tool that gets the bank statement or the PDF file as an argument, and then gives you also a JSON document. And actually it's now not a single JSON document, but JSON lines, so multiple documents, but the general idea stays the same. And if you now run that, and leave everything else untouched, you get a speedup of a factor of 11. So just by replacing this one little step, and not like, I would say it's not overly complicated to do it, but you get a nice, decent speedup here. There are a couple of caveats here. On the one hand, I definitely have to mention pdf extract is definitely less mature than pdfminer6. So for example in this bank of mine there's like this initial page where you have this like hello mr poem and so on for example this is not passed at all this is completely lost. So whether this is applicable it clearly depends on on the pdfs you're using and also maybe on the library more maturing. At the same time also the speedup you get really depends on the pdfs you're using. In this case it's a very simple pdf but in practice I actually observed a factor of 60 so when you had something like takes like hours you talk about minutes and this is like really a decent speed up so suddenly we didn't need to parallelize because you could just use a single core on a laptop and do our data processing and so it depends always a little bit on your use case but in particular for PDFs I would really recommend you to check that out. So let's get to a second step or like a second use case specifically about parsing JSON files. And one thing I will be using here is the arrow file format. I think some people have heard it already. There were some questions also in the lightning talks channel for example and maybe let's take a step back and explain a little bit what the arrow format is. The arrow format is not a data format for like a file format but rather it's a way of arranging data in memory so essentially it's a specification how you structure your bits and bytes such that different languages can use the same data and if you look for example on the official documentation you will find bindings for C, C++ and so on but also Python and Rust which are the two libraries that we are interested in and what is also really nice about Arrow it does not only support primitive types so something like strings, integers, floats, but you can really encode lists, you can encode structures, maps, you can even encode unions in this data format. So in principle, I would say everything you can put down in Python, you can probably also put down in Arrow. Whether you can then use it efficiently in the library, that's a different case, but at least you can put it into this binary format in the first step. And what also is really nice is, if you're going the PyO3 route, there's really high-level bindings also for that in the arrow package itself, so it's also really easy to interface it and really write efficient functions that take arrow data frames in and then give out arrow data frames out. Maybe to give you a little bit of a feeling how it looks like in practice, so consider you have two columns, one with nullable integers, so some items are missing, and then maybe you have some nullable text column. let's say, hello world, with a missing value in between. And what Arrow does is essentially it specifies how you structure simpler arrays from it. So, for example, for the integer array, you would essentially encode one array of 64-bit integers and then a second array with a bit mask where essentially the missing values are. And then for the strings, you would encode everything as UTF-8 characters and then include essentially another array that gives you the offsets where the different strings start and again a nullability mask. And now the trick is, because everybody knows how to encode strings, all the different libraries know how to encode strings, you can now use the same data in pandas, polars, HuggingFace is using it for the HuggingFace datasets, I think PyTorch is now using Arrow underneath, and obviously also in your own tools you can use it. And then essentially write Arrow and then put it into all these other tools. So the use case here is essentially parving like a heap of JSON documents, convert it into like an arrow file, and then do some analysis on top of it in Python. The data set in question is a Smiljan playlist Spotify data set. It's something like 30 gigabytes uncompressed, 5.4 gigabytes compressed, and has essentially listings of different playlists with artists and so on, and then you can ask what is the most popular artist in this dataset, for example. And when we compress it down into this arrow file format, this is actually like some intermediate format that is more used for internal exchange, but you can also use it with all the tools that speak arrow typically, then we will compress it down to 2.3 gigabytes, which is actually a nice dataset that you can either load in RAM on a normal desktop, or even you just use a memory map and then use efficiently without loading it at all. So when we do it in Python, we, again, this route of defining types and doing this type definition way, we have the schema of the data set. As I said before, you have essentially the different playlist names, and then you have different columns. In particular, you have this tracks column, and this is actually what I mentioned before. You can encode really these complex data types. So in this case, for example, we have this list of dictionaries, and in the dictionaries and you have artist names and track URIs and other things. So how do we then convert it in Python to this arrow format? It looks a bit unwieldy but I just saw that I have enough time to walk you through it slowly. So first of all you need to load individual JSON documents. I didn't decompress the zip file but just use it as is. And essentially load the individual JSON blocks I think something like 1,000 playlists in a file and then 1,000 files in the overall zip file. Then there's a processing step which we need to do because dates in Arrow are actually encoded in milliseconds, not in seconds. So we need to multiply by a factor of 1,000. And then finally, we iterate over the different columns and build out the individual arrays. Maybe, sorry, in Arrow speak, a table is a data frame, roughly. and an array is a series roughly if you're coming from pandas and one thing you may already notice here that there are a couple of loops in in python so first of all you have this loop over the playlist and then you have this loop over the different columns first of all but that's not a slow loop but you also have this loop over the individual playlists so essentially you loop over all the items i mentioned that already here because surprisingly this is very fast which Which is a shame for me later on. On the Rust side of things, we use the same strategy as before. We define our types, then use code generation for the types. And so all we now need to do is we need to load the JSON files and then convert everything back to an arrow file. One thing I'm using here is a library I've written, but there are other libraries around, in particular the main arrow library has now this feature also built in. But the main idea is to use the same code that we generated for serialization and deserialization to also build our arrays that we can then essentially use in our analytics workflows. So, in principle, this is all the code of the package, of this program. There's a little bit more, but in principle, that's it. The operations are more or less the same. First of all, we loop over the individual files and we load the JSON document. The specifics are here arrow-specific, but in principle it's the same step, bus-specific, sorry, but in principle it's the same step as in Python. Then we iterate over the individual playlists and perform this factor of thousand conversion and then essentially we push it into this builder object to build our arrow arrays in RAM. And then we can write it to disk and then use it for something like in polars, for example, to do the analysis on top. So let's compare compare the different results. And as I said, this will look a little bit disappointing in the beginning, but we'll get to it. This is a little bit more exciting, so bear with me. So first of all, you see there's not that much performance difference from Python to Rust, and then also I plucked out this part which builds the arrays manually, added something like 200 lines more of code to build the arrays manually, but also that didn't really help that much. So if you go from Python to the Rust implementation that I've shown you before, we have something like a factor of 10%. And actually, the only reason why it's even faster is because we have this multiplication by a thousand, to be very honest. So if you have obviously more transformation, it becomes slower in Python. If you then build the arrays manually, you have something like a 60% speedup, so it gets a little bit more exciting. However, I was kind of lying to you, because the big problem here is that the Python implementation actually doesn't really help or doesn't work at all because I have something like 16 gigabytes of RAM but the overall data set is something like 32 gigabytes right and so I can't fit it in RAM and I can't really process it and the reason here is or the issue here is that you have very many repeated elements so one artist is appearing in very many playlists and so clearly this can be very much compressed down and something that you can use here is a technique called dictionary encoding in arrow or categorical encoding in pandas to really essentially replace all these repeated elements. And this is already done by the Rust side but not done by the Python side. If you then implement that in Python as well, it suddenly becomes very slow, not very slow but much slower than before. And suddenly you get a decent speedup with Rust. So you have something like 2.4 if you are using the lazy route and 3.3 if you build the arrays manually. And this is then really this overhead where you have then this custom string manipulations where you do dictionary lookups and suddenly Python becomes slow. One thing that's also maybe of interest here is if you just do the JSON parting in Rust and throw away all the data, you just go over the data set, pass the files, and then count the playlists or whatever. So something very simple. This is more or less the whole runtime of the program. And this is actually why Python is also so fast. Because in principle, everything you're doing in the simple implementation is you're loading the file from disk, you're doing the JSON parsing, all of that is done in C, and yeah, this is efficient. So essentially whether you use Rust for that or Python doesn't really make a difference. I think you really get a benefit here suddenly when you do something that is a little bit more complicated and at your own custom processing steps. And then suddenly you see decent speedups without doing anything overly complicated. So let me conclude. I tried to show you a little bit how to use Rust for data processing, at least like a first glimpse of that. In principle, Rust is fast and memory efficient all over the box. That's not also always the case. There are some simple ways to shoot yourself in the foot. But overall, the naive implementation is already fast quite often. The big caveat here, I think, the Rust libraries, in particular for data processing, are not as high quality as the Python ones, so you should expect to write more code yourself. For example, for the PDF thingy, if we wanted to have the first page of the bank statements, probably we would need to have to write our own extraction library on top of the low-level PDF library, so clearly that's more complicated. But in general, I think everything should be more or less there and hopefully also quality improves over time. I talked a bit about different strategies. I think if you want to go for some high quality extension module that you want to share with the community, you should definitely go with PyO3 or some other related binding package. I would still pick PyO3 probably. But what I hoped also to show you is that quite often you can also get away with something hacky. little command line tool, don't care about all these particularities of the C API, just do something JSON in, JSON out, or write a file to disk and then use it in your workflow. The general strategy always is essentially use type hints as often as possible or type definitions in Rust and essentially generate most of the code, so it's really also easy to get started and efficient to get going. And then if you're in this like more small data regime, so for the PDF case, we have essentially medium-large files, but then we like extract a couple of words. So this is like something where the encoding doesn't need to be efficient. We can get away with using JSON. Something like the Spotify playlist data set where we have very many rows, then suddenly it makes sense to go in this direction of using Arrow for that. And finally, some tips when to incorporate Rust. So first of all I think performance benefits are not always clear cut. So the story of I will add rust and suddenly it's a factor of 10 can sometimes be the case but also quite often it will not be the case. And the underlying reason is essentially that lots of Python libraries are already written in C and are fast out of the box. Here the thing is really if you essentially leave the confines of these libraries and you need to do something custom. I think then I would investigate Rust and look into that as an option. In particular, whenever you need to bring something that is like in some weird format, I would count PDF as a weird format, into something that is Python compatible, I would definitely look into Rust and the more processing steps you have in their data pipeline. So even this little factor of a thousand is already enough, then suddenly doing something in Rust becomes viable. But if you just do load data, put it into another format, I think probably Ross won't help you here. Thank you very much, and I would be happy for any questions.

Speaker 2 [29:44]

Thank you, Christopher, for your talk. Unfortunately, we are out of time for the questions. Sorry, maybe it was misleading about the indications. But if you have your questions, we have them on Slido. I recommend you to talk to Christopher directly for your questions. And thank you for attending. A round of applause again for Christopher.

Christopher Prohm

Christopher is a data scientist and long-time Python user. Recently he started using Rust for data projects and became interested in how to combine both languages.

Social card for talk: Pragmatic ways of using Rust in your data project