No-Code Change in Your Python UDF for Arrow Optimization

Key Takeaways

I'm Yakin Guan. I'm the the software and staff software engineering data bricks uh up a software foundation member and a spot uh PMC member and committer and I'm the top contributor uh in Apach spot ripple. All right. So today I'm going to talk about the the no code change in your Python UDF for arrow optimization. I made the title very interesting, exciting. So I'm happy that the my talk is accepted. Hopefully the contents are exciting too. Before all I would like to celebrate myself because I hit the 2000 commit in the apostas repo. There's a bit of a uh like rough moments like there are some people who helped me a lot for the decade contribution. So I'd like to say thanks to them. All right, back to the topic introduction. Um like uh this session is actually designed for the beginners. So I just want to quickly check this guy like this anyone who doesn't know Pepar the Python UDF. Okay. Okay. Okay. Okay. Okie dokie. Yeah. So this is the um basically what I'm going to talk today. So um I'm going to go through the what's python udf. The python udf you know you can run some random projects right. So you can um directly import the third party libraries and then you can like use like you can reuse the function that you used to running in the like single machine. You can use that function to the the distributed machine. Um however the the performance itself is not very good. I I'll discuss why it's slow and then like talk about it a little bit later. So um we introduced the the arrow optimized Python UDM. So um it um it optimized the the Apache arrow. Do you guys know Apach arrow? No. Yes. Yes. Yes. Yes. Yes. I'm going to briefly go through that too. Um so you can turn it on and off. There's a like new parameter introduced. So you can change your code or you can set the one configuration. So like out of the your um workflow you can just switch the spot configuration on. Then you can leverage the like up to two times faster Python UDF without changing anything. Of course, I will um talk about the benchmark wizard too. Um uh as I just mentioned, no need to change. And then lastly, I will talk a little bit about the data bricks stuff because like I'm the datab bricks employee and they like pay me. So I will just talk a little bit about it at the last. All right. So Python UDF just in case you guys are not very familiar with it. I mean seems like most of people know. So I'll quickly go through. So what's Python UDF? So Python is user defined function. So um uh you can write the your own logic in Python and then you can run it in the the pipark. So for instance on the on your right side there's a um this my upper that say um use your Python function and then you know that's basically making your string uppercase right. So this is actually very widely used. uh we analyzed the the notebook usage in the pip database before and then 95% of users are using the python udf when they are using python notebook um so um we um actually uh investigated why it's so popular because we have it like different many APIs right we have RDD API map part like map in pandas and pandas udf um it basically enables um directly uh your existing workload to be distributed right suppose that the you have a python function that's running in local computation and then you want to distribute the computation you can just wrap it with the decor the decorator the UDF then you can directly um use that to the in the distributed environment so it needs they minimize change and then um there are some like libraries that are not supported in the the PIP park right in like third party libraries the Python ecosystem all the dependencies you can directly use with that the Python UDF plus there are some logics that um cannot be uh written easily in the pispark expressions. I intentionally made up on your right there's a um uh there's a the python udf that basically just filter out the special characters and then um just extract the username. This kind of expressions are well maybe you are able to do it with the python expressions but you're like let's suppose that you're just python user you're very familiar with the python um logics then it's like pretty difficult to do it. So because of the both reason the python udfs are like very popular as I mentioned but the performance is not that great. Um there are several reasons behind but but basically you're like serializing the row by row one by one and then you're using pickle that is not designed for the efficient um um encoding but like as you guys might already know we introduced the pandas udf right so the basically um it's sending batch by batch it's using apach arrow format I will talk about it later um so it sending the batch by batch and then when you execute them then uh it it can be executed with the vectorzed oper uh operation that's leveraging the you know the the the modern the CPU design I'll talk about it later so why don't we use the the pandas udf pandas is hard it's difficult to learn to be honest um it's a little bit embarrassing to say but uh like my first time I started to learn pandas that's when I implemented the pandas udf so that was the like a few years ago and at the time what I learned that you know Firstly, the indexing and slicings are hard. I mean, I didn't even know like spark data frame doesn't have index. So, I didn't know what's index. I mean, why is it there? I mean, why do you need it? And then slicing and then you like has some like brackets and then what's something weird inside and then name and then I didn't know what's even index at the time. And then the pandas itself is that I'm a fan of pandas by the way, you know, I need to do like introduce my, you know, feature. So, the lots of implicit behaviors like for instance here, you know, df.call call and then I did the the the the comparison with a right I was expecting that there are singular singular value like true or false but suddenly like it it returns a like series of like bunch of multiple like booleans so this kind of things like n to deny like python users it's pretty difficult maybe like if you're data scientist I mean you might think that I'm dumb but you know it was like hard for me and then there's no um like uh to the na knife python noogers um they're not really familiar with the the vectorzed operations like I mean I was just like I would say you want to change something in the list I would do the for loop but like in like pandas way you have to do like a vectorzed way that's why um we introduce a arrow optimized python udf so you don't need to change your existing function but you can leverage the arrow optimization under the hood by um switching on and off configuration. All right. So bottlenecks. So uh I just uh mentioned that the the regular pipe and UDF it's great. It's widely used but there are a bit of it's not quite fast, right? The major reason is that um we are doing one by one sd the world by role and then we using pickle. Pickle is not designed for the you know optimized like transferring stuff. you know in fact when you when you pickle for instance a string then the size even become bigger. So um we are switching that the pickle to the apach arrow under the hood. So I mean since you guys know I will just briefly quickly go through the apach arrow is that the memory format. So it's optimized memory format. So you can put your data uh in the very efficient uh that like small size. So you can read and write quickly uh easily and then it supports it supports the a lot of the different um languages for instance R or Python, Java. So it's very um efficient way to uh transfer data between process or even within process or within um uh the shared memory. All right. So P pandas udfs are already using the arrow right. So when they send the batch to the the pipeline um worker then um they are already using the arrows. Then the idea is that the why don't we just uh keep the existing trans transferring logic but we just convert the arrow batch back to the the roles and then we just apply the function for individual roles. So that was the idea. Yeah. So uh this is the the feature that I added. So we have a two ways to enable this. So one is that that you we have a new parameter called use called use arrow. So you can specify that to um true then you can control the the individual udf to use the arrow or not or you can just leave your workload as is and then you can uh enable the configuration called the spark.sql.execution the python udfer enabled. Yeah. So you can um specify the configuration in your like spark like config file or you can just fix your workload to be the enable and then test it out. This is actually introduced from the spark 3.5 but um like there has to be some optimization done. So this is actually still in progress. So um I would expect this to be ready in spark 4.1 and then we actually even want to enable this by default. So if you are interested you can just try it out and see if it works for DBR. uh the data bricks runtime is supported from 14.0. All right. So by using the arrow under the hood actually there's another benefit. Um so the first one the top one is that we are using arrow the second one is a disabled arrow you you guys might know that the when the return returned instance is different from the specified SQL type. For instance, here the select I'm um here the return type I'm specifying the return type as string, right? But the data frame it has a dates. So when it when it happens when arrow is enabled, it it returns a relatively um reasonable string representation of dates, right? But when you use the Python UDF and when the types are different, then it shows something that weird strings. The reason behind is that the uh we used the library that reads the pickled instance to the the Java instance within the JVM. So here the second uh we read the Java the pickled instance to the the Gregorian calendar that instance and then that's converted to the string. So my point is that that when you use arrow then you you you can get the more reasonable um uh the the representation of string types or other casted types. All right. So benchmark before I talk about the benchmark measured in the arrow optimized um um Python UDF. No no this is not a this is pandas pandas pandas udf this is different one. So uh this is the the pandas udf benchmark. So when we did that it shows that the performance uh improvement up to 100 times faster. It's pretty high right. So to be honest I was expecting that such kind of performance improvement but when we did the performance um benchmark then it shows up to the only two times I mean even fast which is great but we need to figure out why right so we investigated why. So the pandas is like you know back backed by the numpy array right and then the numpy array itself the like support the vectorzed operation. So which is like a CPU friendly the same right. So basically like numpy allows you your CPU to just bring the data batch and then just apply the uh computation in batch. But basically the pandas is like uh implemented by under the hood by the numpy. when you execute the function then uh you can do the vector operation but since in this uh python arrow optimized python udf we are changing the arrow batch to the individual row right and then we apply function for row by row so you cannot really do the vectorzed operation so this almost like a little bit um critical in the performance wise but still since we switch the serial serial format from the pickle to arrow that boosted the performance up to two times. Oopsie. Yeah. And then you know for the data bricks little bit only one slide. Um we are trying to do do you guys know Photon? I don't know Photon. Raise your hands. Okay. So Photon is basically like native native engine. So basically the the Spark engine is written by C. So like they're using like vector operation under the hood. So before that uh we didn't support the Python UDF with photon but since we're you know switching the the the from pico to the arrow batch right so basically we can address the the the batch together and we can combine that to boton so we have the prototype now and then it's working in progress so the python udf will be photonized and then uh we are like exciting and then we're benchmarking now okie dokie so key takeaways so forget about everything I said I you know just you know just this is summary yeah so summary so UDFs are widely used but slow right because people can run these some custom logics and then they can use the third party libraries they can use some like a logics that cannot be written in the pispark expressions so we introduced the arrow optimized python udf that can be easily turned on and off we added the one parameter at udf and then we added the one uh spark configuration and then we ran some benchmark Spark. Um, this is the not fast as the pandas UDF but reasonably fast because it can be up to like two times faster. And then the one very good one benefit is that the you don't have to change your code. So you just try to configuration on and see if it becomes faster or not. So it's pretty easy to try. Lastly, data bricks specific. We're trying to botonize it. So if you're like datab bricks customers, then you might be more happier. Okie dokie. Yeah, that's the what I prepared today, right? Thank you so much.