Skip to content

Commit e6ce5a9

Browse files
authored
Merge pull request #646 from novatechflow/feature/spark-dataframes
Feature/spark dataframes
2 parents 07cc9fd + a473152 commit e6ce5a9

25 files changed

Lines changed: 1348 additions & 33 deletions

README.md

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,20 @@ Apache Wayang provides a flexible architecture which enables easy addition of ne
7373

7474
For a quick guide on how to run WordCount see [here](guides/tutorial.md).
7575

76+
### Spark Dataset / DataFrame pipelines
77+
78+
Wayang’s Spark platform can now execute end-to-end pipelines on Spark `Dataset[Row]` (aka DataFrames). This is particularly useful when working with lakehouse-style storage (Parquet/Delta) or when you want to plug Spark ML stages into a Wayang plan without repeatedly falling back to RDDs.
79+
80+
To build a Dataset-backed pipeline:
81+
82+
1. **Use the Dataset-aware plan builder APIs.**
83+
- `PlanBuilder.readParquet(..., preferDataset = true)` (or `JavaPlanBuilder.readParquet(..., ..., true)`) reads Parquet files directly into a Dataset channel.
84+
- `DataQuanta.writeParquet(..., preferDataset = true)` writes a Dataset channel without converting it back to an RDD.
85+
2. **Keep operators dataset-compatible.** Most operators continue to work unchanged; if an operator explicitly prefers RDDs, Wayang will insert the necessary conversions automatically (at an additional cost). Custom operators can expose `DatasetChannel` descriptors to stay in the dataframe world.
86+
3. **Let the optimizer do the rest.** The optimizer now assigns a higher cost to Dataset↔RDD conversions, so once you opt into Dataset sources/sinks the plan will stay in Dataset form by default.
87+
88+
No extra flags are required—just opt into the Dataset-based APIs where you want dataframe semantics. If you see unexpected conversions in your execution plan, check that the upstream/downstream operators you use can consume `DatasetChannel`s; otherwise Wayang will insert a conversion operator for you.
89+
7690
## Quick Guide for Developing with Wayang
7791

7892
For a quick guide on how to use Wayang in your Java/Scala project see [here](guides/develop-with-Wayang.md).

guides/spark-datasets.md

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
<!--
2+
3+
Licensed to the Apache Software Foundation (ASF) under one or more
4+
contributor license agreements. See the NOTICE file distributed with
5+
this work for additional information regarding copyright ownership.
6+
The ASF licenses this file to You under the Apache License, Version 2.0
7+
(the "License"); you may not use this file except in compliance with
8+
the License. You may obtain a copy of the License at
9+
10+
http://www.apache.org/licenses/LICENSE-2.0
11+
12+
Unless required by applicable law or agreed to in writing, software
13+
distributed under the License is distributed on an "AS IS" BASIS,
14+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
See the License for the specific language governing permissions and
16+
limitations under the License.
17+
18+
-->
19+
20+
---
21+
title: Spark Dataset pipelines
22+
description: How to build Wayang jobs that stay on Spark Datasets/DataFrames from source to sink.
23+
---
24+
25+
Wayang’s Spark backend can now run entire pipelines on Spark `Dataset[Row]` (a.k.a. DataFrames). Use this mode when you ingest from lakehouse formats (Parquet/Delta), interoperate with Spark ML stages, or simply prefer schema-aware processing. This guide explains how to opt in.
26+
27+
## When to use Dataset channels
28+
29+
- **Lakehouse storage:** Reading Parquet/Delta directly into datasets avoids repeated schema inference and keeps Spark’s optimized Parquet reader in play.
30+
- **Spark ML:** Our ML operators already convert RDDs into DataFrames internally. Feeding them a dataset channel skips that conversion and preserves column names.
31+
- **Federated pipelines:** You can mix dataset-backed stages on Spark with other platforms; Wayang will insert conversions only when strictly necessary.
32+
33+
## Enable Dataset sources and sinks
34+
35+
1. **Plan builder APIs:**
36+
- `PlanBuilder.readParquet(..., preferDataset = true)` (Scala) or `JavaPlanBuilder.readParquet(..., ..., true)` loads Parquet files into a `DatasetChannel` instead of an `RddChannel`.
37+
- `DataQuanta.writeParquet(..., preferDataset = true)` writes a dataset back to Parquet without converting to RDD first.
38+
2. **Prefer dataset-friendly operators:** Most unary/binary operators accept either channel type, but custom operators can advertise dataset descriptors explicitly. See `DatasetChannel` in `wayang-platforms/wayang-spark` for details.
39+
3. **Let the optimizer keep it:** The optimizer now assigns costs to Dataset↔RDD conversions, so once your plan starts with a dataset channel it will stay in dataset form unless an operator demands an RDD.
40+
41+
## Mixing with RDD operators
42+
43+
If a stage only supports RDDs, Wayang inserts conversion operators automatically:
44+
45+
- `SparkRddToDatasetOperator` converts an RDD of `org.apache.wayang.basic.data.Record` into a Spark `Dataset[Row]` (using sampled schema inference or `RecordType`).
46+
- `SparkDatasetToRddOperator` turns a `Dataset[Row]` back into a JavaRDD.`Record`.
47+
48+
Both conversions carry non-trivial load profiles. You’ll see them in plan explanations if you mix dataset- and RDD-only operators.
49+
50+
## Developer checklist
51+
52+
- **Use `RecordType` when possible.** Providing field names in your logical operators helps the converter derive a precise schema.
53+
- **Re-use `sparkExecutor.ss`.** When writing custom Spark operators that build DataFrames, use the provided `SparkExecutor` instead of `SparkSession.builder()` to avoid extra contexts.
54+
- **Watch plan explanations.** Run `PlanBuilder.buildAndExplain(true)` to verify whether conversions are inserted. If they are, consider adding dataset descriptors to your operators.
55+
56+
## Current limitations
57+
58+
- Only Parquet sources/sinks expose dataset-specific APIs today. Text/Object sources still produce RDD channels.
59+
- ML4All pipelines currently emit plain `double[]`/`Double` RDDs. They still benefit from the internal DataFrame conversions but do not expose dataset channels yet.
60+
61+
Contributions to widen dataset support (e.g., dataset-aware `map`/`filter` or ML4All stages) are welcome.

wayang-api/wayang-api-scala-java/src/main/scala/org/apache/wayang/api/DataQuanta.scala

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ import org.apache.wayang.core.optimizer.costs.LoadProfileEstimator
3636
import org.apache.wayang.core.plan.wayangplan._
3737
import org.apache.wayang.core.platform.Platform
3838
import org.apache.wayang.core.util.{Tuple => WayangTuple}
39-
import org.apache.wayang.basic.data.{Tuple2 => WayangTuple2}
39+
import org.apache.wayang.basic.data.{Record, Tuple2 => WayangTuple2}
4040
import org.apache.wayang.basic.model.{DLModel, LogisticRegressionModel,DecisionTreeRegressionModel};
4141
import org.apache.wayang.commons.util.profiledb.model.Experiment
4242
import com.google.protobuf.ByteString;
@@ -1027,6 +1027,11 @@ class DataQuanta[Out: ClassTag](val operator: ElementaryOperator, outputIndex: I
10271027
writeTextFileJava(url, toSerializableFunction(formatterUdf), udfLoad)
10281028
}
10291029

1030+
def writeParquet(url: String,
1031+
overwrite: Boolean = false,
1032+
preferDataset: Boolean = false)(implicit ev: Out =:= Record): Unit =
1033+
writeParquetJava(url, overwrite, preferDataset)
1034+
10301035
/**
10311036
* Write the data quanta in this instance to a text file. Triggers execution.
10321037
*
@@ -1090,6 +1095,16 @@ class DataQuanta[Out: ClassTag](val operator: ElementaryOperator, outputIndex: I
10901095
this.planBuilder.sinks.clear()
10911096
}
10921097

1098+
private def writeParquetJava(url: String, overwrite: Boolean, preferDataset: Boolean)(implicit ev: Out =:= Record): Unit = {
1099+
val _ = ev
1100+
val sink = new ParquetSink(url, overwrite, preferDataset)
1101+
sink.setName(s"Write parquet $url")
1102+
this.connectTo(sink, 0)
1103+
this.planBuilder.sinks += sink
1104+
this.planBuilder.buildAndExecute()
1105+
this.planBuilder.sinks.clear()
1106+
}
1107+
10931108
/**
10941109
* Write the data quanta in this instance to a Object file. Triggers execution.
10951110
*

wayang-api/wayang-api-scala-java/src/main/scala/org/apache/wayang/api/JavaPlanBuilder.scala

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -63,14 +63,21 @@ class JavaPlanBuilder(wayangCtx: WayangContext, jobName: String) {
6363
createSourceBuilder(new TextFileSource(url))(ClassTag(classOf[String]))
6464

6565
/**
66-
* Read a parquet file and provide it as a dataset of [[Record]]s.
66+
* Read a parquet file and optionally keep it backed by Spark Datasets.
6767
*
6868
* @param url the URL of the Parquet file
6969
* @param projection the projection, if any
70+
* @param preferDataset when {@code true}, emit a Dataset-backed channel
7071
* @return [[DataQuantaBuilder]] for the file
7172
*/
72-
def readParquet(url: String, projection: Array[String] = null): UnarySourceDataQuantaBuilder[UnarySourceDataQuantaBuilder[_, Record], Record] =
73-
createSourceBuilder(ParquetSource.create(url, projection))(ClassTag(classOf[Record]))
73+
def readParquet(url: String,
74+
projection: Array[String]): UnarySourceDataQuantaBuilder[UnarySourceDataQuantaBuilder[_, Record], Record] =
75+
readParquet(url, projection, preferDataset = false)
76+
77+
def readParquet(url: String,
78+
projection: Array[String] = null,
79+
preferDataset: Boolean = false): UnarySourceDataQuantaBuilder[UnarySourceDataQuantaBuilder[_, Record], Record] =
80+
createSourceBuilder(ParquetSource.create(url, projection).preferDatasetOutput(preferDataset))(ClassTag(classOf[Record]))
7481

7582
/**
7683
* Read a text file from a Google Cloud Storage bucket and provide it as a dataset of [[String]]s, one per line.

wayang-api/wayang-api-scala-java/src/main/scala/org/apache/wayang/api/PlanBuilder.scala

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -136,9 +136,13 @@ class PlanBuilder(private[api] val wayangContext: WayangContext, private var job
136136
*
137137
* @param url the URL of the Parquet file
138138
* @param projection the projection, if any
139+
* @param preferDataset when {@code true}, keep the resulting channel backed by Spark Datasets
139140
* @return [[DataQuanta]] of [[Record]]s representing the file
140141
*/
141-
def readParquet(url: String, projection: Array[String] = null): DataQuanta[Record] = load(ParquetSource.create(url, projection))
142+
def readParquet(url: String,
143+
projection: Array[String] = null,
144+
preferDataset: Boolean = false): DataQuanta[Record] =
145+
load(ParquetSource.create(url, projection).preferDatasetOutput(preferDataset))
142146

143147
/**
144148
* Read a text file from a Google Cloud Storage bucket and provide it as a dataset of [[String]]s, one per line.
Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing, software
13+
* distributed under the License is distributed on an "AS IS" BASIS,
14+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
19+
package org.apache.wayang.basic.operators;
20+
21+
import org.apache.wayang.basic.data.Record;
22+
import org.apache.wayang.core.plan.wayangplan.UnarySink;
23+
import org.apache.wayang.core.types.DataSetType;
24+
25+
/**
26+
* Logical operator that writes {@link Record}s into a Parquet file.
27+
*/
28+
public class ParquetSink extends UnarySink<Record> {
29+
30+
private final String outputUrl;
31+
32+
private final boolean isOverwrite;
33+
34+
private final boolean preferDataset;
35+
36+
public ParquetSink(String outputUrl, boolean isOverwrite, boolean preferDataset, DataSetType<Record> type) {
37+
super(type);
38+
this.outputUrl = outputUrl;
39+
this.isOverwrite = isOverwrite;
40+
this.preferDataset = preferDataset;
41+
}
42+
43+
public ParquetSink(String outputUrl, boolean isOverwrite, boolean preferDataset) {
44+
this(outputUrl, isOverwrite, preferDataset, DataSetType.createDefault(Record.class));
45+
}
46+
47+
public String getOutputUrl() {
48+
return this.outputUrl;
49+
}
50+
51+
public boolean isOverwrite() {
52+
return this.isOverwrite;
53+
}
54+
55+
public boolean prefersDataset() {
56+
return this.preferDataset;
57+
}
58+
}

wayang-commons/wayang-basic/src/main/java/org/apache/wayang/basic/operators/ParquetSource.java

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,8 @@ public class ParquetSource extends UnarySource<Record> {
5757

5858
private MessageType schema;
5959

60+
private boolean preferDatasetOutput = false;
61+
6062
/**
6163
* Creates a new instance.
6264
*
@@ -124,6 +126,16 @@ public ParquetSource(ParquetSource that) {
124126
this.projection = that.getProjection();
125127
this.metadata = that.getMetadata();
126128
this.schema = that.getSchema();
129+
this.preferDatasetOutput = that.preferDatasetOutput;
130+
}
131+
132+
public ParquetSource preferDatasetOutput(boolean preferDataset) {
133+
this.preferDatasetOutput = preferDataset;
134+
return this;
135+
}
136+
137+
public boolean isDatasetOutputPreferred() {
138+
return this.preferDatasetOutput;
127139
}
128140

129141
@Override

wayang-platforms/wayang-spark/src/main/java/org/apache/wayang/spark/channels/ChannelConversions.java

Lines changed: 50 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,18 +19,23 @@
1919
package org.apache.wayang.spark.channels;
2020

2121
import org.apache.wayang.basic.channels.FileChannel;
22+
import org.apache.wayang.basic.data.Record;
2223
import org.apache.wayang.basic.data.Tuple2;
24+
import org.apache.wayang.core.api.Configuration;
2325
import org.apache.wayang.core.optimizer.channels.ChannelConversion;
2426
import org.apache.wayang.core.optimizer.channels.DefaultChannelConversion;
27+
import org.apache.wayang.core.plan.executionplan.Channel;
2528
import org.apache.wayang.core.types.DataSetType;
2629
import org.apache.wayang.java.channels.CollectionChannel;
2730
import org.apache.wayang.java.platform.JavaPlatform;
2831
import org.apache.wayang.spark.operators.SparkBroadcastOperator;
2932
import org.apache.wayang.spark.operators.SparkCacheOperator;
3033
import org.apache.wayang.spark.operators.SparkCollectOperator;
3134
import org.apache.wayang.spark.operators.SparkCollectionSource;
35+
import org.apache.wayang.spark.operators.SparkDatasetToRddOperator;
3236
import org.apache.wayang.spark.operators.SparkObjectFileSink;
3337
import org.apache.wayang.spark.operators.SparkObjectFileSource;
38+
import org.apache.wayang.spark.operators.SparkRddToDatasetOperator;
3439
import org.apache.wayang.spark.operators.SparkTsvFileSink;
3540
import org.apache.wayang.spark.operators.SparkTsvFileSource;
3641

@@ -108,6 +113,32 @@ public class ChannelConversions {
108113
() -> new SparkObjectFileSource<>(DataSetType.createDefault(Void.class))
109114
);
110115

116+
public static final ChannelConversion DATASET_TO_UNCACHED_RDD = new DefaultChannelConversion(
117+
DatasetChannel.UNCACHED_DESCRIPTOR,
118+
RddChannel.UNCACHED_DESCRIPTOR,
119+
() -> new SparkDatasetToRddOperator()
120+
);
121+
122+
public static final ChannelConversion CACHED_DATASET_TO_UNCACHED_RDD = new DefaultChannelConversion(
123+
DatasetChannel.CACHED_DESCRIPTOR,
124+
RddChannel.UNCACHED_DESCRIPTOR,
125+
() -> new SparkDatasetToRddOperator()
126+
);
127+
128+
public static final ChannelConversion UNCACHED_RDD_TO_UNCACHED_DATASET = new DefaultChannelConversion(
129+
RddChannel.UNCACHED_DESCRIPTOR,
130+
DatasetChannel.UNCACHED_DESCRIPTOR,
131+
ChannelConversions::createRddToDatasetOperator,
132+
"via SparkRddToDatasetOperator"
133+
);
134+
135+
public static final ChannelConversion CACHED_RDD_TO_UNCACHED_DATASET = new DefaultChannelConversion(
136+
RddChannel.CACHED_DESCRIPTOR,
137+
DatasetChannel.UNCACHED_DESCRIPTOR,
138+
ChannelConversions::createRddToDatasetOperator,
139+
"via SparkRddToDatasetOperator"
140+
);
141+
111142
public static Collection<ChannelConversion> ALL = Arrays.asList(
112143
UNCACHED_RDD_TO_CACHED_RDD,
113144
COLLECTION_TO_BROADCAST,
@@ -119,6 +150,24 @@ public class ChannelConversions {
119150
HDFS_OBJECT_FILE_TO_UNCACHED_RDD,
120151
// HDFS_TSV_TO_UNCACHED_RDD,
121152
CACHED_RDD_TO_HDFS_TSV,
122-
UNCACHED_RDD_TO_HDFS_TSV
153+
UNCACHED_RDD_TO_HDFS_TSV,
154+
DATASET_TO_UNCACHED_RDD,
155+
CACHED_DATASET_TO_UNCACHED_RDD,
156+
UNCACHED_RDD_TO_UNCACHED_DATASET,
157+
CACHED_RDD_TO_UNCACHED_DATASET
123158
);
159+
160+
private static SparkRddToDatasetOperator createRddToDatasetOperator(Channel sourceChannel,
161+
Configuration configuration) {
162+
DataSetType<Record> type = DataSetType.createDefault(Record.class);
163+
if (sourceChannel != null) {
164+
DataSetType<?> sourceType = sourceChannel.getDataSetType();
165+
if (Record.class.isAssignableFrom(sourceType.getDataUnitType().getTypeClass())) {
166+
@SuppressWarnings("unchecked")
167+
DataSetType<Record> casted = (DataSetType<Record>) sourceType;
168+
type = casted;
169+
}
170+
}
171+
return new SparkRddToDatasetOperator(type);
172+
}
124173
}

0 commit comments

Comments
 (0)