From Row-Wise to Columnar: Speeding Up PySpark UDFs with Arrow and Polars

Speaker 1 [00:00]

attendees, welcome from my side. I will raise awareness to the questionnaire platform, talks.python.com.de. There you can ask the questions for the session. Now I want to introduce our amazing speaker, Emilius from Frontier Media's SA. He works as a software engineer He has been working as a dev engineer at CERN and he will present his amazing talk about from row-wise to columnar, speeding up PySparks, UDFs, with IRO and Polaris. So this is your stage, Emilio, here you go.

Speaker 2 [00:46]

Thank you very much. Good morning from my side. My name is Emilios. For most of the people it's difficult to pronounce my last name, that's usual. I have a long name. So today here I am going to present you a problem that I had currently in my work. And I work as a software engineer for a company which is an open access publisher, which means that we publish research papers in our journals. And I work mainly on the data product. We use Spark and I found some complications. I tried some stuff and I wanted to present them to you. So when we have somebody new into Spark the idea is that we do big data, distribute computing and so on and the first advice that we give is okay whenever you want to do something try to use native Spark functions if you cannot you can write your own functions which we call them UDFs or user defined functions. What they do is that the main principle is that I want to do a task which is not possible natively. You can model the data with them, perform transformations, and usually the tip we give is you should use them with caution. Why that? Because they are helpful, you are doing your task, but in the end they slow you down. And this is something true, but I will try to change today the aspect a bit. So here I present you the data flow over a Spark UDF, and what we have, we have a partition with safe rows into the Java virtual machine. This thing needs to be serialised, arrow needs to be serialised and converted to bytes, passes through a socket to move from the JVM to the Python worker, and the Python worker reads those bytes, it deserialises them, it creates the Python object, you pass this Python object to the function that you have written, you get the result, and you do the opposite trip to go back to the JVM. What is the problem here? Actually we have two problems. The first problem is that if you take a look at this graph or if you listen to me describing it, I say serialise, serialise, serialise, serialise. So we are doing four times this action and this is a CPU bound action and a very heavy one and we don't want it. We are adding serialisation overhead to what we are doing. Plus, we are executing the computation row by row. Here I want to note something important is that many people don't think when they do the computation what kind of library they use. What I mean by that? Using a Python library will be slow. If you can find a library that does the same thing natively, it will be much better even if you use a Spark UDF, and we will see it later, I don't have included the example, but I have examples where I have used a very specific library, because the Python version was much slower. The evolution of UDFs in Spark is that we have what we call Spark UDF, which is the a classic one. Everybody starts by this UDF. It was introduced with Spark 0.7. As we said, the serialisation is being done by PIC and it's row by row. Then we move at some point to Spark 2.3, which is Pandas UDF. Here things are changing a bit. The serialisation becomes comes arrow, we are executing by series, the engine is mostly numpy I or C. We continue further with arrow optimised UDFs, which is something different of a concept. It's mostly an improvement over the Spark UDFs. What improvement? The serialisation here uses arrow again, but the execution is row by row. And we come to today, which is arrow UDFs. It's a completely different aspect because you use arrow as the serialisation and the execution is vectorised and mostly using C++ kernels. Following I will continue with the Pandas UDF workflow. Here we have a batch of rows, around 10,000, 10,000 is the default batch record for Spark. And towards the serialisation part, the JVM, what it does, it changes the rows to columns. So they can enter the ROA APC, and then this ROA APC converts to bytes, they pass through the socket. When they pass through the socket, they get deserialised and they become a record patch. And this record patch is converted to Pandas, to Pandas series, so you can pass it through the UDF. You get the result, and you have again the round trip back to the JVM. A few notes here. Of course, we have avoided with Arrow the row-by-row deserialisation. But it's not enough. Why it's not enough? We have high memory usage. The conversion from Arrow to Pandas and Pandas to Arrow will give you... You will keep the objects two times in memory. So this is high memory usage. And the Pandas data frame from some data frame operations is single-threaded, which is slowing us down. So I talked about a new concept which I did not introduce, so I will do it now. And this is Apache Arrow. It's an in-memory columnar format. And what it achieves? It solves the copy and convert bottleneck. And what is a copy and convert bottleneck? It's what I previously talked. In theory, we keep, and in practice, we keep two times the same data in memory. Why we do that? Because PySpark is a wrapper over Spark, Spark works over the JVM, PySpark works over a Python process. So these processes, without Arrow, they need to have a copy of the data. With Arrow, they don't need that, you don't need to copy the data, you can have the data available to both systems, and why is that? Because you are reading directly from the bytes, you are reading the memory. And this is what we call zero copy. Arrow also helps in another aspect, which is a bit of lower level and goes to the hardware. So I said before that we are doing some CPU operations, all of them, actually. And the biggest problem is that how we retrieve the data, how we use the data to execute it on the CPU. And the CPU has caches, L1, L2, L3. And when you are on the row layout, so you are not having a column, you are having the whole row being serialized into the cache. But what happens if you need only one of the columns of this row? You are filling the cache with data that you don't need, that you are going to request, but you are not going to request. So the next data that you will need, it will not be in the cache, it won't fit. So you will get always cache misses. This is changing with a columnar format. Why? Because you are just getting the row, you are applying it continuously in the cache, and you're not having misses. And the hardware prefacer helps in that, in the sense that it predicts always the code that you, the data for the code that you will need. Continuing on the next UDF is the arrow-optimized UDF. As I told you, it's an improvement in one aspect over the Spark UDF. The improvement is that you see that the first part, until the record batch, does the same. Rows to columns, columns to bytes, over the socket, and you get the record batch. But the biggest issue here is unpack. So because you haven't solved the row-by-row execution, even if you have a batch of columns, you need a specific row to pass inside your function. and you do the opposite trick. So the Python overhead remains. And we go towards ROUDF. ROUDF follows the same pattern of conversion passing through the socket and so on, but we don't have any conversion. As you see, from the record batch, we have a zero copy, And we go inside the function, which is using PyArrow compute and performance computation. So the data stays columnar, and it's also computed with C++ kernels. A small issue, or a bigger one, is that it requires PyArrow compute, which is more of a low-level API. And if you want to perform multi-step transformations, then it's a bit harder, let's say, to use from what we have been used to with DataFrame APIs. And what is the solution to that? The solution to that is to use another DataFrame library as your query engine, actually, which is Polars. Polars is quite a new library. I think it started in 2020. It is a very fast library, DataFrame library. It's written in Rust. It has multi-threaded query engine, so we bypass the problem that we have with Pandas and single-threaded. And it's built on Apache Arrow memory layout, so it means that it already supports Apache Arrow. And it has a user-friendly and expressive DataFrame API. Now, besides the evolution of UDFs, we have better integration with Arrow API in Spark. And what we have? From Spark 3.3, we have mapping Arrow, which is the analogous to mapping Pandas, but we don't have the conversion to Pandas. From Spark 4, we have applying Arrow, which is the analogous to applying Pandas. And we don't, again, have the conversion. And we can convert a Spark data frame to arrow from arrow. And if we want to see how mapping arrow data flow works, it's the same idea what we have with arrow UDF, but you have an iterator of record batches. So the idea is the same. There is no conversion to pandas, and you always stay on the arrow format. Actually, the longer you stay on the arrow format, the better it is for your computations. Now, these are the methods that I have tested. So I started with a Spark UDF. I picked the PIC, which is the classic approach and the arrow transport, Pandas UDF, Arrow UDF, and Mapping Arrow. And for Arrow UDF and Mapping Arrow, I chose two query engines, PyArrow and Polars. These are the tests that I'm going to show you. These are tests, actually, that existed in my code. And I started fixing them. The history is that I started from HTML cleaning. And at some point, I got a call from a colleague telling me that this was really slow. I took a look at the beginning, and it was written in beautiful soup, so it was slow, without using a LXLM parser. Then what I did to improve it is just I used LXLM. I didn't even use even a better libraries, but LXLM is much faster. And then I tried to follow the same idea for other transformations. So in theory, all of my tests, these five tests broken, actually for me it's, let's say, one transformation, one transformation in the data frame. And I will start with the first one. I hope it's readable. I think it is. So I am going to show you, for each of the tests, the implementation. As you will see, we will go through it. But for more or less, for more of them, the code is the same. So the structure is the same. It's only the execution of the specific test method that we are using. So in this case, we have the classic UDF. I think everybody can recognize. Then we go to the classic UDF with the arrow transportation. You just see that the change is only just a flag. And we go to Pandas UDF. We have two options here, actually. We have, let's say, the vectorized in quotes option, which is str-normalize, and using Unicode data library, which is a Python library, with the apply. And you call it with column, which is the standard way of calling it. Moving on, you will see, first of all, that we can do the same calling with ROUDF, no matter of the implementation. On top, you see the PyArrow implementation. We are reading an array, a PyArray compute array, actually batches of this array, and we are returning an arrow array. On the second implementation, same idea. We need to do it in polars. Here, you see that for this example, I can read directly the column as polar series, which is arrow interoperability with polars. And I do the computation, which is the same thing like what I did with arrow, what I did with the rest of the Unicode data or strNormalize, and I save the result to arrow. And the difference, mapInArrow, mapInArrow gives you a bit more control, actually. So if you want to call it, you do dataFrame.mapInArrow, the function that you want to call, and you need to provide the output schema. Without it, it won't run. The idea as a computation is the same. What changes is that you process in batches. A note here that when I first tried this way, I passed the whole data frame, which is of course wrong, because the computation increases significantly. So let's see the results on this. The results are this, and now I will try to remember what the next slide writes by heart. The first thing that I observe is that for the Spark UDFs, I gain around 10 per cent, so the arrow transport versus pick gives me around 10 per cent, three seconds. The second thing I see is that Pandas UDF and Pandas UDF vectorised are almost equal, And this doesn't make sense. Actually it makes, because if you read the code, you will see that str normalised from pandas is not vectorised, so it's calling the same thing like the Unicode data library. And then we have map in arrow, which is faster over polars. And this is, sorry, map in arrow which is faster over ROUDF. And this is due to the fact that there is some overhead, UDF overhead, that you need to do. So you take the column, you pass it to the function, and you return it back. With map in arrow, you are providing the schema. So you're having the schema, you're doing a transformation. And as you see, the gap is really small. So this is a result on a single column. And this is what I told you. And also, okay, for this case, PyArrow and Polar are within the same approach. Now, I did the same thing, but this time I passed more columns. I passed four columns. I won't show you the code, because it's a for loop and it's the same for this case. But I will go straight to the results. And the results are this. So here the results become more interesting. What I observe, the gap between Spark UDF and Spark UDF are optimized, increased, and that's normal. It's something that I expected. Vectorized here wins because we have four columns, so the classic approach with .apply has more overhead. You can see that for mapping arrow, pi arrow and polars have a small gap difference of around three seconds. And this is happening because, and this is happening also in arrow UDF, and this is happening because we have the conversion gap. Although the read is zero copy, then we have a conversion gap. And that the gap between the fastest approach to the slowest approach increased. So in this case, I will proceed with the next test, which is a normalisation, which I remove also a specific Unicode group. And this is happening like that. So in the beginning, I keep the same structure, I just added that I don't want this specific category of Unicode characters. You will see here that I don't have vectorize because it uses the RE module and the RE module of Python doesn't support this kind of operations. Moving to the arrow UDFs, you see that we added a new line. That's what I mean by low-level API for PyArrow. I mean, it's a bit more difficult and less expressive, at least for me, from what pollers provide. For me, what pollers provide makes more sense to what I have been used to in Spark. And of course, the same idea comes to mapping arrow, that nothing changes, we work in batches and we perform the same computation. And the results are these. So what we see here is that, of course, native code wins. That's the most visible aspect. And wins because it's what I told you before. But besides thinking what transportation you have or what is the Python for loop overhead, you need to think how you write the actual computation. The gap increases a lot, more than the previous case. Still we have the same amount for our transport, and we expect that the work that needs to to be done in the side of Python is more. So this is for, let's say, string normalization. And I go back to the problem of HTML cleaning, which was the beginning of my idea. Here I used, actually, the best library that I know. I changed to Lexbor. And this is the method that we use for cleaning. Nothing fancy. I mean, it's a very easy method to write. And we start calling this method. Again, we do it with Spark UDF, Spark with UDF are optimized, and Pandas UDF. We follow the same approach for RO UDFs. And the thing here to note is that in this case, there is no native kernel, neither in PyArrow nor in Polars. So what we do is that we call the main function that does the transformation, and we convert the column to a list. So inside here, do you believe that this is going to run fast? I mean, it's still inside the Python loop. And the same idea for mapping arrow, only this time, as always with mapping arrow, we do it with batches. So the benchmarks are these. And the result becomes interesting in the sense that from a first point of view, you You see an outlier, we'll see why. And all of them are being quite close. So literally you could have used anything you wanted of this. Even a Spark UDF wouldn't have any, let's say, disadvantage over the other UDFs or Mapinaro. Why this is happening? This is happening because here, the transport part is much lower over the computational part. Actually, where we spend the time here is mostly on the computation part. Now, why Spark UDFR optimized is that much slower and is an outlier? This is happening because here we have variable length strings. We are making a batch, and we are trying to pick from this batch. This is simply not efficient in this case. It is efficient, for example, for small strings like names and so on, but for variable length strings that can be very, very long, it's totally inefficient. And you see here, and you see also, sorry, that in difference with the previous example, The compute part is heavier. In the normalization part, the transport part was a big aspect because the computational part was lightweight. And I move also to, let's say, string similarity. Also here, I tried to use the best library. Actually, I use cydflib and not the dflib from Python. Actually, I think that when I tested the DiffLib, it was two times slower. So by changing only the library, I gained a benefit of 100%. What I do is that let's say that I want to compare two full names. The idea is that I will compare the first with the second, and then the second with the first, And from their scores, I will get the mean of it. How I do that? I take the first name or the second. I split it. And for the parts that I have split, I do my computation. I sequence match works with longest common substring, which is a dynamic programming problem. And this is the heavy part here. that we need somehow to make this part faster. But is it possible to make a dynamic programming problem faster? That's a good question, and we will see. So again, we follow the same approach. In this case, we just pass, let's say, a struct or zip the names. Nothing really different from the previous case, because we have a function, and we call this function. We do the same in PyArrow and Polars. And in this case, in Polars, we pass a struct, and we map it with the elements. And we do the same also for mapping arrow. Now, going over the results, what do we expect? I assume we expect something like that. So we have come towards the problem that we have previously, and the problem that we have previously is that I have a heavy computational problem, so no matter what I use, the result is always going to be the same. But there is a but here. And the but is what got me thinking is I want to make this faster. Can I do it? And the answer is yes. So I changed a bit the approach. And by changing the approach, I mean the following. Instead of doing the stuff with a for loop and passing it to sequence matcher, I said, if I want this to be vectorised, I will do the following. I will explode all my columns. I will run them. I will do, OK, everything that I do in the previous, in the main method. I will run it for one name and for the other two times. I will join the results, and I will run the longest common substring. I will join the results, and I will get it back. The good thing of that is that Polar has a plug-in, which is called PDS, and has this functionality. And it's actually quite optimised also for space. And now, I can even write this natively, but OK. But I can use ROUDF. And the results become something like that. So, when I go into the vectorised space, I get around 2.2 performance increase. Amazing, eh? For something that does an explode in the group by, I mean, explode in Spark, well, it's not happening in Spark here, it's happening in Polars, would be quite computationally costly. So why this is happening? This is happening for two reasons. The first reason is that the heavy computational part is the part that it was optimised by running it into native code, in this case Rust, which is using Rayon, parallel threading and so on and so on. The second part is that how explodeby works in Arrow and in Polars. And the idea is something like that. Let's say that we have this table. And arrow works in buffers. And what I mean by buffers, they have a buffer which is offset and a buffer which is data. For integers, you don't have this problem. You have just the values as a buffer. But for strings, let's say, you have the offset, the buffer, offset, and the data. What that means is that the offset points to you from where the string starts and where it ends. And the list, of course, is a nested structure of the same idea. So what is really happening before the explode is the diagram that you see there. You have the two columns, row ID and name. And then, for the rest, which is a list, you have the offsets which are pointing really the strings which are contingent in memory. So when you explode, what you do is in the end, you just allocate some new rows for the other columns, but for the nested structure, you are using again the child array that you have, and it's just a reference count, it's not a copy. So this gives you the opportunity not to lose time with an explode, not to be computationally pick. And, of course, you drop the offset buffer, because then you don't have a list, you don't need it, you don't have something nested, you have something flat. If we check the same for Spark, the idea of the exploding is the same, but here you carry with you, it's the same that I was trying to explain you in CAS locality, is that you carry the whole row with you. So you cannot say that, oh, from this row, I want this part. So if you want to explode it, you will explode everything. You will create many, many objects of the originals. So, takeaways. When your function is compute heavy, and it's a Python call per row, every UDF type converges, and I think this is clear from my results. The real win is if you can avoid the Python loop, so if you can find a way to do it natively, either in C, C++, or Rust. That's very good. And you need to match the tool to the problem you are trying to solve. And this was my presentation. Thank you very much for your attention.

Speaker 1 [33:38]

Also from my side, thank you for this amazing talk. We have questions, and people can also still use talks.byton.de to ask questions. The first question will be a very long question with two sub-questions. The first question of the sub-question is, how much this approach bridges the gap between natures, spark functions, and UDFs? And the second of those, did you try to write a native Spark function as an UDF and compare it with a native Spark function?

Speaker 2 [34:15]

So, for the second, I have not done it to run a Spark native function as a Spark UDF function. But there are presentations from the Apache Spark community that they have done that and they show real improvements. And for the first one, also, it's something that I have not tested. I'm not sure if he means actually to write specific natively in Java or Scala and do the comparison, but I have not done it also.

Speaker 1 [34:47]

Okay, I see. It was quite an early question, so it might have also been answered by your talk, because you really benchmarked it. It was really nice. Sorry. Does it make sense to use parallel polars execution, given the fact that Spark already executes one executed process per physical core?

Speaker 2 [35:13]

I am not sure I get it.

Speaker 1 [35:16]

Should I...

Speaker 2 [35:17]

Should I repeat it? Yes.

Speaker 1 [35:18]

Does it make sense to use parallel polos execution? Given the fact that Spark already executes one executor process per physical core.

Speaker 2 [35:32]

Okay, it depends. So, Polars is multi-threaded on a single node. Spark distributes the workload. So, for example, in this case, the benchmarks that I have tested, I have done them in a single node of eight cores. I could have done them in almost everything. For the stuff that they are, for the last example, which is native, If I start adding more nodes, of course, I will parallelize more the problem. So I expect even more to drop. But I wanted to keep it stable, the benchmark, to be accurate. But yes, you can parallelize. But the thing is that Polar parallelizes only in a single node. It's not distributed. For distributed, they have Polar's cloud. So in this case, what you do, you use either mapping arrow or arrow UDF to distribute the problem to all your worker nodes.

Speaker 1 [36:36]

Thanks, I think this answers the question. So, like, there's no question left anymore. So I think you will be around here still in the venue the full day. If you have questions or worries, like later on, just approach him or connect with him on LinkedIn and you are always open to answer the questions.

Speaker 2 [36:56]

questions.

Speaker 1 [36:56]

Yes. Super. This was the last talk before the lunch break and now I want to give you an applause and then have a nice lunch break.

Speaker 2 [36:57]

Yes.