How to create a custom Spark Encoder in ... java

🕥 13 min

What is a Spark Encoder ?

An Encoder is a wrapper class that specifies how to serialize and deserialize data with the Spark Structured Streaming framework.

In this framework, datasets are typed and they are schema aware so every Spark transformation called on a dataset needs to provide an Encoder for its output type so that the framework can know how to serialize and deserialize its content. For example, the signature of the map transform is as this:

Why this blog post ?

Because Spark Structured Streaming Encoders are complex to write, have no java API and they are not very documented so I thought that what I learnt writing them could be useful to the community.

Catalyst

To understand Encoders we need to focus a bit on Catalyst. Spark Structured Streaming uses Catalyst optimizer to optimize execution plans. And Encoders are part of the Catalyst plan. Catalyst sees a pipeline as a tree made of TreeNodes. Here is the Catalyst workflow:

The entry point of Catalyst is either a SQL abstract syntax tree (AST) returned by the SQL parser if the user has used direct Spark SQL or a Dataframe object if the user has used the Dataframe/Dataset API. The result is an unresolved logical plan with unbound attributes and data types. During the analysis phase, Catalyst does catalog lookups and attributes mapping to fill these placeholders and ends up with a logical plan. This plan is then optimized by applying standard rule-based optimizations to end up with the optimized logical plan.  Then Catalyst uses rules to turn the optimized logical plan into one or more physical plans with physical operators (for example dataset.map) that match the Spark execution engine. It then selects a plan using a cost model that consists of applying cost evaluation rules to the physical plans.

When we come back to the subject is with the last phase: the code generation. Among the Treenodes types, there are Catalyst Expressions and custom Encoders are Catalayst Expressions called ExpressionEncoders. Catalyst Expressions contain java code strings that get composed to form the code AST of the pipeline which is then compiled to java bytecode using Janino compiler. This bytecode is the one that is executed when the pipeline runs. Let's see how to write ExpressionEncoders:

Custom Encoder (ExpressionEncoder)

To manage serialization in Spark Structured Streaming framework, you could use the Encoders available in the Encoders utility class. They can manage serialization of primitive types or beans with either java or kryo serialization. But let's say you want to write a custom Encoder because, for example, you develop a framework based on Spark and you want to allow your users to provide serialization code. As an example, let's take Apache Beam which lets the user define his serialization code in the Coder class.


Spark Encoders have no java api, so if you want to code in java because your code base is in java, you'll need to use some bindings. Maybe there are better ways, but as I'm not a scala developer, I could be unaware of them :)


As said, a custom Encoder is an ExpressionEncoder

ExpressionEncoder groups serializer and deserializer. You need to specify your serializer and deserializer as Expression instances. Let's see the serializer part:

Serializer

Things to point out in that code:

• It is an Encoder so it has no SQL representation (remember, Catalyst Expressions are broader than Encoders) so we  implement NonSQLExpression and it has one input and one output so we extend UnaryExpression.
• There are several methods to override:
• child(): the Encoder is part of Catalyst tree (see above) so we need to keep track of its child (the input of the UnaryExpression)
• doGenCode(): this method is responsible for producing the java code strings that will be compiled by Janino.
• datatype(): that is the type of the result of the evaluation of the Expression. In our case binary (because we are serializing an object to binary).
• the other overrides productElement(), productArity(), canEqual() and consequently equals() and hashcode() are due to the fact that there is no Java API of ExpressionEncoder so we need to implement Scala product specifics.

Code generation

Let's focus on the interesting method doGenCode():

This method generates the code in the comment line 32 in the form of java strings.

It returns an ExprCode, see line 52. This Block is constructed through string interpolation. This Block creation is managed line 49.

Now that the global architecture of this method is clearer, there is some pieces that could look weird:

• To access an object that is not a local variable part of the generated code block, we need to add a Catalyst reference to it, see line 24. In our case we reference the coder which contains the user provided serialization code.

• As said above the serializer is a UnaryExpression. This Expression has only one input Expression which is its child (see lines 25 and 32). We need to concatenate the child code and the actual serialization code so that everything can be compiled by Janino line 52.

Instantiate the serializer

Here is the code to create the EncodeUsingBeamCoder object (the serializer part of our ExpressionEncoder). To instantiate this class we need to pass it a reference to its child in the Catalyst tree (remember, the input Expression of the UnaryExpression). To obtain a reference to the Catalyst input Expression, we do like this: BoundReference(0, new ObjectType(clazz), true). There is only one input (because EncodeUsingBeamCoder is a UnaryExpression) so we get it at index 0 and we indicate the Datatype of the input and its nullability.

Spark physical plan

The serialization part of the physical Catalyst plan comes like this

SerializeFromObject [encodeusingbeamcoder(input[0, org.apache.beam.sdk.util.WindowedValue, true], WindowedValue$FullWindowedValueCoder(VarIntCoder,GlobalWindow$Coder)) AS binaryStructField#11]

WindowedValue is the type to serialize from  (the clazz)  and WindowedValue$FullWindowedValueCoder(VarIntCoder,GlobalWindow$Coder) is the Beam Coder provided in the pipeline.

Deserializer

The deserializer class DecodeUsingBeamCoder that you can see in the full code link below is completely symetric to the serializer class EncodeUsingBeamCoder. The only thing worth mentioning is its instanciation:

Instantiate the deserializer

Here is the code to create the DecodeUsingBeamCoder object (the deserializer part of our ExpressionEncoder). Here also, to instantiate this class, we need to pass it a reference to its child in the catalyst tree (remember, the input Expression of the UnaryExpression). To obtain a reference to the catalyst node, we do like this: new Cast(new GetColumnByOrdinal(0, BinaryType), BinaryType).

Here again we get the first input Expression at index 0 (cf UnaryExpression), it is of type BinaryType because we are deserializing bytes. And the Cast allows Spark to treat the Expression more efficiently as it allows Catalyst to treat it as binary.

Spark physical plan

The deserialization part of the physical Catalyst plan comes like this

DeserializeToObject decodeusingbeamcoder(cast(binaryStructField#4 as binary), org.apache.beam.sdk.util.WindowedValue, WindowedValue$FullWindowedValueCoder(VarIntCoder,GlobalWindow$Coder)), obj#6: org.apache.beam.sdk.util.WindowedValue

WindowedValue is the type to deserialize to (the clazz) and WindowedValue$FullWindowedValueCoder(VarIntCoder,GlobalWindow$Coder) is the Beam Coder provided in the pipeline

Performances

Even though Janino compiler is fast, compiling java strings to bytecode takes time. I measured a big performance gain when I reduced the size of the generated code and replaced it by as much compiled code as possible. This is why in full code link, encode() method in the serializer and decode() method in the deserializer are compiled code and not string code inside doGenCode() method. Another gain to this approach is to enable debugging of this compiled part of the code.

Full code link

EncoderHelpers class in the Apache Beam project

Understand Apache Beam runners: focus on the Spark runner

🕥 5 min.

Previously on Apache Beam runners 😀

In the previous article, we had a brief overview of what an Apache Beam runner is. This article will dig into more details.

The previous article introduced this very simple pipeline:

We saw that the Beam SDK translates this pipeline into a DAG representing the pipeline in the form of Beam transform nodes. Now let's see how a Beam runner translates this DAG. Let's say that the user choses Apache Spark as the target Big Data platform when he launches his pipeline.

The runner (at last)

The job of the runner is to translate the pipeline DAG into a native pipeline code for the targeted Big Data platform. It is this native code that will be executed by a Big Data cluster. If Spark is chosen, then the runner translates this DAG below into the Spark native pipeline code below. For the sake of simplicity the Spark pipeline is pseudo-code.

Composite and primitive transforms: the level of translation

In the previous article we talked about Beam transforms (primitive transforms and composite transforms). In the continuation of the blog, we will refer to "composite transform" as just "composite" and "primitive transform" as just "primitive".

The DAG above is the expanded DAG: on the left hand-side are the Beam transforms of the user pipeline. But among these transforms only Read is a primitive. The others are implemented by Beam SDK as composites of other primitives. Composites can also be made of composites themselves (like Count transform is made of Combine transform) but in the end they are always made of primitives (Pardo or GroupByKey). And these primitives are what the runner translates.

But in some cases, the runner can chose to translate at a composite level of the graph, not at a primitive level depending on the target Big Data technology capabilities. Indeed, if there is a direct correspondance of the Beam composite in the target API, the runner does not decompose the composite into its primitives and translates directly the composite. The green boxes in the DAG represent the level of translation. In our example, there is a direct equivalent of Beam Combine composite to a Spark Aggregator (agg in the Spark pipeline).

The translation itself

The translation occurs when the pipeline is run (pipeline.run() is executed). To translate the DAG, the runner visits the graph. All Beam runners work the same, only the target API changes with the chosen runner.

The first step is to detect the translation mode (batch or streaming) by searching for Beam BoundedSource (like Elasticsearch for example) or UnBoundedSource (like Kafka for example). Knowing the translation mode, the runner can chose

• the proper Spark DataSourceV2 to instantiate either implementing ReadSupport (batch) or MicroBatchReadSupport (streaming)
• the proper Spark action to execute on the output dataset either foreach (batch) or writeStream (streaming)

Then the DAG visit continues and each node is translated to the target native API : Read gets translated to a Spark DataSourceV2 that creates the input dataset, Pardo gets translated to a Spark flatmap that is applied to the input dataset and so on until the output dataset.

Pipeline run

When the visit of the DAG is done, the runner applies the action chosen above to the output dataset to run the pipeline. At this point the spark framework executes the resulting native Spark pipeline.

Introduction to Apache Beam and the runners

🕥 7 min.

This is a first blog to introduce Apache Beam before going into more details about Beam runners in the next blogs.

What is Apache Beam ?

To define Apache Beam let's start with a quote

«Batch / streaming? Never heard of either.»
(Batch is nearly always part of higher-level streaming) 

Beam is a programming model that allows to create big data pipelines. Its particularity is that its API is unified between batch and streaming because the only difference between a batch pipeline and a streaming pipeline is that batch pipelines deal with finite data whereas streaming pipelines deal with infinite data. So, batch can be seen as a sub-part of the streaming problem. And Beam provides windowing features that divide infinite data into finite chunks.

Let's take a look at the Beam stack:

The pipeline user code is written using one of the several language SDKs (java, python or go) that Beam provides to the user. The SDK is a set of libraries of transforms and IOs that allow him to input data into his pipeline, do some computation on this data and then output to a given destination. Once the pipeline is written, the user choses a Big Data execution engine such as Apache Flink, Apache Spark or others to run his pipeline. When the pipeline is run, the runner first translates it to native code depending on the chosen Big Data engine. It is this native code that is executed by the Big Data cluster.

Primitive transforms

The Beam SDK contains a set of transforms that are the building blocks of the pipelines the user writes. There are only 3 primitives :

Pardo: It is the good old flatmap that allows to process a collection element per element in parallel and apply them a function called DoFn

GroupByKey: This one groups the elements that have a common key. The groups can then be processed in parallel by next transforms downstream in the pipeline.

Read: The read transform is the way to input data into the pipeline by reading an external source which can be a batch source (like a database) or a streaming (continuously growing) source (such as a kafka topic).

Composite transforms

Composite transforms: All the other transforms available in the SDK are actually implemented as composites of the 3 previous ones.

As an example, the Reduce of Beam which is called Combine and which allows to do a computation on data spread across the cluster is implemented like this:


Other examples of composite transforms provided by the SDK are FlatMapElements, MapElements or Count that we will see in next chapter. But the user can also create his own composites and composites of composites.

A simple Beam pipeline

Let's take a look at a first simple Beam pipeline:


This is the usual Hello World type big data pipeline that counts the occurrences of the different words in a text file. It reads a text file from google storage and the result of the count is output to google storage. This is a very simple straight pipeline. Not all the pipelines have to be straight that way, it is there to serve as a baseline example for the continuation of the blog and to illustrate some key concepts of Beam :

• Pipeline: the user interaction object.
• PCollection: Beam abstraction of the collection of elements spread across a cluster.
• TextIO: this IO is used to read and write from/to the text file. Reading part of the IO is in reality a Read transform and writing part of the IO is in reality a Write transform
• Count: Combine transform that counts occurrences
• FlatMapElements and MapElements: They are composite transforms of ParDo that are there for convenience.

The resulting DAG

When Beam executes the above pipeline code, the SDK first creates the adjacent graph to represent the pipeline. It is known as the DAG (Direct Acyclic Graph). For each transform of the pipeline code, a node of the DAG is created.  

•  Read node corresponds to the Read transform of TextIO (the input of the piepeline)
• Write node corresponds to the Write transform of TextIO (the output of the pipeline)
• All the other nodes are a direct representation of the transforms of the pipeline.

Please note that only the Read transform is a primitive transform as described in the above paragraph, all the others are composite transforms.

But what about the runner ?

This is when the runner enters the scene. The job of the runner is simply to translate the pipeline DAG into a native pipeline code for the targeted Big Data platform. It is this code that will be executed by a Big Data cluster such as Apache Spark or Apache Flink.

You want to know more about the runner ? The next article describes what the runner is and uses the Spark runner as an example