Repository navigation
[SPARK-14083][WIP] Basic bytecode analyzer to speed up Datasets - #24515
aokolnychyi wants to merge 2 commits into
Conversation
|
@dbtsai and I collaborated on this. The approach we take was originally described by @JoshRosen. The implementation and scope are very different, though. As the PR description says, the goal is to start a discussion and not to suggest a properly working optimal approach. To the best of my knowledge, there are multiple companies that are trying/tried to implement this functionality. It would be great if everyone can share their progress/ideas/pain points so that we can, hopefully, improve typed Datasets in Spark. @rednaxelafx @kiszk @maropu @rxin @dongjoon-hyun @cloud-fan @gatorsmile |
|
Test build #105086 has finished for PR 24515 at commit
|
|
We will make changes to Spark to open up some refection and |
|
@aokolnychyi Thanks for your contribution! @rednaxelafx will help review this PR in the next few days. He is the expert in this area |
|
Test build #105111 has finished for PR 24515 at commit
|
|
Thanks for your work, @aokolnychyi and @dbtsai ! I've worked on a continuation of @JoshRosen 's prototype about two years ago, so I have some first-hand experience on both the implementation details and the applicability of this direction. In the meantime, though, I'd really like to call on the community to share their use cases of using the existing typed Dataset operations, so that we'll be able to better evaluate how much benefit will this project bring to real-world queries. In a lot of cases, simply moving uses of the typed Dataset operations to the equivalent untyped DataFrame operations can substantially speed up queries; third-party solutions like Quill can provide a typed API for Scala but directly generate untyped operations that are fast to begin with (in Quill's case, the generated code is SQL). So for people that are not content with the current performance of the typed Dataset operations, they already have two directions they can pursue:
Or, if a query uses bulk operations like There are a few cases where users may be forced to use the typed Dataset operations, e.g. when they need to use Structured Streaming APIs like So questions to everybody:
Thanks! |
|
I will start reviewing it next week when I come bacj to Japan. I have the same questions as the 2nd and 3rd of @rednaxelafx |
|
Thanks for your work! @rednaxelafx talked about several questions that I think most people are interested to know. For simple typed Dataset operations, it is quite easily to use untyped DataFrame operations instead. But for complex typed Dataset operations, I think it might be more challenging to write uses of the typed Dataset operations to the equivalent untyped DataFrame operations. However, it might be where bytecode analyzer can be really valuable as it automatically converts the operations written in host language to untyped DataFrame operations. So I'm interested to know that what kind of code can be safely analyzed for speeding up, and what kind of code mightn't be. |
|
As the person that initially filed the ticket, I actually no longer believe in it, for the following reasons:
|
|
In some APIs there's no choice but leverages typed API, and for structured streaming (flat)MapGroupsWithState does, which is recommended on all stateful use cases except basic cases. Would it be better if we present untyped API as well for APIs which are only available in typed? |
|
I like the direction suggested by @rednaxelafx. That’s exactly what we planned to do next: identify concrete use cases of the typed API and see how we can optimize them (e.g. bytecode analysis, straightforward equivalent untyped operations, extending the untyped API, etc). I believe this discussion should be driven by use cases. @viirya the question of what operations can be safely optimized is mostly open. I think it is possible to support closures without side-effects (e.g. writing to files/logging) that operate on primitives, boxed values, case classes/POJOs. Potentially, we can handle side-effects via @rxin I know at least four companies that tried something around this topic. This PR is an attempt to bring everyone to the same table and have a consensus on the feasibility of bytecode analysis and future optimizations of the typed API. We are all aligned that there should be enough use cases that benefit from this and that the implementation must be reliable. Because of the complexity and maintenance overhead, I don't consider this to be in Spark directly. At most, a separate package. |
|
We also have a reasonable collections of Dataset API use cases at LinkedIn, especially centered around offline feature engineering pipelines which relies on complex transformation logics that are not straightforward to express using DataFrame operations. We are also working on a similar prototype to address Dataset performance issue. We are trying to find a balance between bringing the benefits of bytecode analysis and dealing with its complexity. Instead of trying to fully convert the lambda function into Catalyst expression, which might run into many corner cases, we are rather focused on identifying which fields of the domain objects are being accessed and leverage that information in column pruning optimization to cut down serde and IO overhead. Want to chime into this discussion and provide our 2 cents and work with the community to see how we can push this Dataset performance enhancement effort forward. |
|
I like the approach @Victsm describes |
|
In addition to the ideas discussed here, I think we should also benchmark the raw constant-factor overheads of UDF / UDAF / typed operations to see whether there's any straightforward optimizations that will speed up existing workloads without the added complexity of bytecode analysis / closure conversion. For example, it looks like there's room for improvement in how we invoke UDFs with primitive input type arguments (https://issues.apache.org/jira/browse/SPARK-27684). Through careful benchmarking we might be able to uncover other low-hanging wins. |
|
thanks for posting this, I haven't looked at the code in detail yet, but we are also interested in this area. In particular we are interested in converting lambdas and udfs into full catalyst expressions. Once you have the catalyst expression you can do more optimizations with it. In our cases we started to look at this for the columnar processing side of things. If you have the catalyst expression then you can map that into a GPU operation. The more time you can keep the data on the GPU, the more performance you can gain. Copying back and forth is inefficient. I think that applies for many types of columnar processing, if you can keep it in columnar without having to switch back and forth to rows, the more benefit you will have. We had a few people start to look at this. Originally they started with javasist but then switched to use JVMCI (but that is only available in jdk > 8 and very specific oracle versions of jdk8). The main reason they switched from javassit is when there are multiple lambdas within the same class they couldn't differentiate them since the lamdba classes generated at runtime can't be relfected or instrumented. If you only have 1 lambda per class it was fine. I'm not sure how many times this will be an issue but thought I would mention it here. I agree with many of your points that need to be discussed and decided upon. I think if we can keep it pluggable like you are proposing people can try different things out. I think one of the main things is to know when to not try to convert or give up. If you do that quickly enough I'm very curious of other people experience here, @rednaxelafx have you had time to write up your thoughts from previous experience? |
|
We're closing this PR because it hasn't been updated in a while. If you'd like to revive this PR, please reopen it! |
Disclaimer! The purpose of this PoC PR is to trigger a discussion on whether it is feasible to leverage bytecode analysis to speed up Datasets. This PR shows what can be achieved by interpreting stack-based bytecode directly. The current version does not handle all edge cases. If the community agrees that bytecode analysis can give substantial benefits, we will need to define the scope (i.e., which cases we are targeting) and decide which approach to take (e.g., stack-based bytecode interpretation, flow-control graphs, AST, Jimple). I would like to emphasize that this PR doesn't answer those questions. Instead, it shows what can be done with the simplest approach. Also, there is no intention to integrate this logic into Spark directly. It might be a separate package.
Scope
Bytecode analysis can be used for several purposes:
Rewriting closures is challenging but gives more benefits. So, this PR focuses on the second use case.
Right now, the code covers typed map and filter operations that involve primitive values, boxed values (not all are implemented), objects that are represented as structs (e.g., case classes). The same logic can be applied to UDFs/UDAFs and other operations.
Algorithm
There are two ways to derive Catalyst expressions:
Both approaches have their own tradeoffs, which are well-described in comments to SPARK-14083. This PR translates stack-based bytecode directly and simulates what happens to the operand stack. The optimizer is supposed to simplify the derived expression. I think it is valid to start with the simplest approach and evaluate its limitations before considering a more advanced implementation.
Originally, we wanted to cover trivial cases. However, the current scope is much bigger, so we should consider an intermediate format and will be glad to hear more opinions. As it was mentioned before, the decision highly depends on our target use cases.
This PR adds a new optimizer rule
RewriteTypedOperationsthat usesTypedOperationsin order to convert typed operations into untyped.TypedOperationsfollows a trivial algorithm that is described below.Step 1: Get closure bytecode
First of all, we need to get bytecode for the closure. Scala 2.12 migrated to Java lambdas and uses
LambdaMetafactory(LMF) (seems like there are rare cases when Scala doesn't use LMF). This PR relies on the existing logic in Spark to obtainSerializedLambda, which has enough information to get bytecode for closures.Scala uses "adapted" methods for LMF to encapsulate differences in boxing semantics (i.e., unboxing null in Scala gives 0 and NPE in Java). The current code will obtain bytecode of the non-adapted method whenever the args are primitives to avoid a round of unnecessary boxing/unboxing.
See
ClosureUtils$getMethodfor more information.Step 2: Build a correct local variable array
Once we have bytecode for our closure, we need to build a local variable array that references correct Catalyst expressions. This can be achieved by translating the deserializer for a typed operation. Deserializers define how data is converted from the Spark internal format into Java objects. Frequently, deserializers contain
StaticInvoke,Invoke,NewInstanceexpressions, which can be translated using the same algorithm.See
TypedOperations$convertDeserializerfor more information.Step 3: Create an operand stack
The next step is to create an operand stack for storing partial Catalyst expressions. The current code uses
mutable.ArrayStack[Expression]for this purpose.Step 4: Interpret instructions
Once we have bytecode, the array of local variables and the operand stack, we can follow our bytecode instructions one by one and simulate what happens to the operand stack.
Step 5: Assign the result back to the expected attributes
Once we have the result expression, we need to assign it to columns. At this point, the serializer is important as it contains information about the expected attributes and their data types.
Trying things out
DatasetBytecodeAnalysisBenchmarkandBytecodeAnalysisSuiteshow currently supported cases. The focus is on the conceptual approach and not on addressing every possible method/instruction.Open Questions