Skip to content

Spark - Accept an output-spec-id that allows writing to a desired partition spec #7120

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
szehon-ho merged 11 commits into from
Apr 13, 2023
Merged

Spark - Accept an output-spec-id that allows writing to a desired partition spec #7120

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
1a4e765
[Spark 3.1] Add option to specify output spec
gustavoatt Mar 15, 2023
afda885
[Spark 3.2] Add option to specify output spec
gustavoatt Mar 15, 2023
d0c5a65
[Spark 3.3] Add option to specify output spec
gustavoatt Mar 15, 2023
df343d9
Compute the outputSpec only during the actual write
gustavoatt Mar 27, 2023
684f1cf
Address comments
gustavoatt Apr 7, 2023
8546b63
Simplify output spec ID precondition check
gustavoatt Apr 11, 2023
7c4235f
Simplify tests
gustavoatt Apr 12, 2023
601c24c
Remove unnecessary comment
gustavoatt Apr 12, 2023
f0eda8d
Move output spec logic into SparkWriteConf
gustavoatt Apr 13, 2023
99529cc
Minimize SparkWrite diffs
gustavoatt Apr 13, 2023
63d6fab
Remove `final` in local variable
gustavoatt Apr 13, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions spark/v3.1/spark/src/main/java/org/apache/iceberg/spark/SparkWriteConf.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -115,6 +115,10 @@ public String wapId() {
return sessionConf.get("spark.wap.id", null);
}

public Integer outputSpecId() {

@aokolnychyi aokolnychyi Apr 13, 2023

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Initially we were not sure whether to make spec IDs public but I also don't see a good alternative.
I am OK with the overall idea of exposing this.

return confParser.intConf().option(SparkWriteOptions.OUTPUT_SPEC_ID).parseOptional();
}

public FileFormat dataFileFormat() {
String valueAsString =
confParser
Expand Down
3 changes: 3 additions & 0 deletions spark/v3.1/spark/src/main/java/org/apache/iceberg/spark/SparkWriteOptions.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,5 +51,8 @@ private SparkWriteOptions() {}
public static final String HANDLE_TIMESTAMP_WITHOUT_TIMEZONE =
"handle-timestamp-without-timezone";

// Output partition spec ID where writes should go.
Comment thread
gustavoatt marked this conversation as resolved.
Outdated
public static final String OUTPUT_SPEC_ID = "output-spec-id";

public static final String OVERWRITE_MODE = "overwrite-mode";
}
31 changes: 26 additions & 5 deletions spark/v3.1/spark/src/main/java/org/apache/iceberg/spark/source/SparkWrite.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,7 @@
import org.apache.iceberg.io.OutputFileFactory;
import org.apache.iceberg.io.PartitioningWriter;
import org.apache.iceberg.io.RollingDataWriter;
import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
import org.apache.iceberg.relocated.com.google.common.collect.Iterables;
import org.apache.iceberg.relocated.com.google.common.collect.Sets;
Expand Down Expand Up @@ -99,6 +100,7 @@ class SparkWrite {
private final String applicationId;
private final boolean wapEnabled;
private final String wapId;
private final int outputSpecId;
private final long targetFileSize;
private final Schema writeSchema;
private final StructType dsSchema;
Expand Down Expand Up @@ -127,6 +129,15 @@ class SparkWrite {
this.dsSchema = dsSchema;
this.extraSnapshotMetadata = writeConf.extraSnapshotMetadata();
this.partitionedFanoutEnabled = writeConf.fanoutWriterEnabled();

Preconditions.checkArgument(
Comment thread
gustavoatt marked this conversation as resolved.
Outdated
writeConf.outputSpecId() == null || table.specs().containsKey(writeConf.outputSpecId()),
"Cannot write to unknown spec: %s",
writeConf.outputSpecId());
this.outputSpecId =
writeConf.outputSpecId() != null
? table.specs().get(writeConf.outputSpecId()).specId()
Comment thread
gustavoatt marked this conversation as resolved.
Outdated
: table.spec().specId();
}

BatchWrite asBatchAppend() {
Expand Down Expand Up @@ -163,7 +174,13 @@ private WriterFactory createWriterFactory() {
Broadcast<Table> tableBroadcast =
sparkContext.broadcast(SerializableTableWithSize.copyOf(table));
return new WriterFactory(
tableBroadcast, format, targetFileSize, writeSchema, dsSchema, partitionedFanoutEnabled);
tableBroadcast,
format,
outputSpecId,
targetFileSize,
writeSchema,
dsSchema,
partitionedFanoutEnabled);
}

private void commitOperation(SnapshotUpdate<?> operation, String description) {
Expand Down Expand Up @@ -558,6 +575,7 @@ DataFile[] files() {
private static class WriterFactory implements DataWriterFactory, StreamingDataWriterFactory {
private final Broadcast<Table> tableBroadcast;
private final FileFormat format;
private final int outputSpecId;
private final long targetFileSize;
private final Schema writeSchema;
private final StructType dsSchema;
Expand All @@ -566,12 +584,14 @@ private static class WriterFactory implements DataWriterFactory, StreamingDataWr
protected WriterFactory(
Broadcast<Table> tableBroadcast,
FileFormat format,
int outputSpecId,
long targetFileSize,
Schema writeSchema,
StructType dsSchema,
boolean partitionedFanoutEnabled) {
this.tableBroadcast = tableBroadcast;
this.format = format;
this.outputSpecId = outputSpecId;
this.targetFileSize = targetFileSize;
this.writeSchema = writeSchema;
this.dsSchema = dsSchema;
Expand All @@ -586,8 +606,8 @@ public DataWriter<InternalRow> createWriter(int partitionId, long taskId) {
@Override
public DataWriter<InternalRow> createWriter(int partitionId, long taskId, long epochId) {
Table table = tableBroadcast.value();
PartitionSpec spec = table.spec();
FileIO io = table.io();
PartitionSpec outputSpec = table.specs().get(outputSpecId);
Comment thread
gustavoatt marked this conversation as resolved.
Outdated

OutputFileFactory fileFactory =
OutputFileFactory.builderFor(table, partitionId, taskId).format(format).build();
Expand All @@ -598,15 +618,16 @@ public DataWriter<InternalRow> createWriter(int partitionId, long taskId, long e
.dataSparkType(dsSchema)
.build();

if (spec.isUnpartitioned()) {
return new UnpartitionedDataWriter(writerFactory, fileFactory, io, spec, targetFileSize);
if (outputSpec.isUnpartitioned()) {
return new UnpartitionedDataWriter(
writerFactory, fileFactory, io, outputSpec, targetFileSize);

} else {
return new PartitionedDataWriter(
writerFactory,
fileFactory,
io,
spec,
outputSpec,
writeSchema,
dsSchema,
targetFileSize,
Expand Down
84 changes: 84 additions & 0 deletions spark/v3.1/spark/src/test/java/org/apache/iceberg/spark/sql/TestPartitionedWrites.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@

import java.util.List;
import java.util.Map;
import org.apache.iceberg.Table;
import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
import org.apache.iceberg.spark.SparkCatalogTestBase;
import org.apache.iceberg.spark.source.SimpleRecord;
Expand Down Expand Up @@ -157,4 +158,87 @@ public void testViewsReturnRecentResults() {
ImmutableList.of(row(1L, "a"), row(1L, "a")),
sql("SELECT * FROM tmp"));
}

@Test
public void testWriteWithOutputSpec() throws NoSuchTableException {
Table table = validationCatalog.loadTable(tableIdent);
// Drop all records in table to have a fresh start.
sql("DELETE FROM %s", tableName);
Comment thread
gustavoatt marked this conversation as resolved.
Outdated
table.refresh();

final int originalSpecId = table.spec().specId();
table.updateSpec().addField("data").commit();

table.refresh();
Comment thread
gustavoatt marked this conversation as resolved.
Outdated
sql("REFRESH TABLE %s", tableName);

// By default, we write to the current spec.
sql("INSERT INTO TABLE %s VALUES (10, 'a')", tableName);
Comment thread
gustavoatt marked this conversation as resolved.
Outdated

List<Object[]> expected = ImmutableList.of(row(10L, "a", table.spec().specId()));
assertEquals(
"Rows must match",
expected,
sql("SELECT id, data, _spec_id FROM %s WHERE id >= 10 ORDER BY id", tableName));

// Output spec ID should be respected when present.
List<SimpleRecord> data =
ImmutableList.of(new SimpleRecord(11, "b"), new SimpleRecord(12, "c"));
spark
.createDataFrame(data, SimpleRecord.class)
.toDF()
.writeTo(tableName)
.option("output-spec-id", Integer.toString(originalSpecId))
.append();

expected =
ImmutableList.of(
row(10L, "a", table.spec().specId()),
row(11L, "b", originalSpecId),
row(12L, "c", originalSpecId));
assertEquals(
"Rows must match",
expected,
sql("SELECT id, data, _spec_id FROM %s WHERE id >= 10 ORDER BY id", tableName));
Comment thread
gustavoatt marked this conversation as resolved.

// Verify that the actual partitions are written with the correct spec ID.
// Two of the partitions should have the original spec ID and one should have the new one.
Dataset<Row> actualPartitionRows =
spark
.read()
.format("iceberg")
.load(tableName + ".partitions")
.select("spec_id", "partition.id_trunc", "partition.data")
.orderBy("spec_id", "partition.id_trunc");

expected =
ImmutableList.of(
row(originalSpecId, 9L, null),
row(originalSpecId, 12L, null),
row(table.spec().specId(), 9L, "a"));
assertEquals(
"There are 3 partitions, one with the original spec ID and two with the new one",
expected,
rowsToJava(actualPartitionRows.collectAsList()));

// Even the default spec ID should be followed when present.
data = ImmutableList.of(new SimpleRecord(13, "d"));
spark
.createDataFrame(data, SimpleRecord.class)
.toDF()
.writeTo(tableName)
.option("output-spec-id", Integer.toString(table.spec().specId()))
.append();

expected =
ImmutableList.of(
row(10L, "a", table.spec().specId()),
row(11L, "b", originalSpecId),
row(12L, "c", originalSpecId),
row(13L, "d", table.spec().specId()));
assertEquals(
"Rows must match",
expected,
sql("SELECT id, data, _spec_id FROM %s WHERE id >= 10 ORDER BY id", tableName));
}
}
4 changes: 4 additions & 0 deletions spark/v3.2/spark/src/main/java/org/apache/iceberg/spark/SparkWriteConf.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -123,6 +123,10 @@ public String wapId() {
return sessionConf.get("spark.wap.id", null);
}

public Integer outputSpecId() {
return confParser.intConf().option(SparkWriteOptions.OUTPUT_SPEC_ID).parseOptional();
}

public boolean mergeSchema() {
return confParser
.booleanConf()
Expand Down
3 changes: 3 additions & 0 deletions spark/v3.2/spark/src/main/java/org/apache/iceberg/spark/SparkWriteOptions.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,9 @@ private SparkWriteOptions() {}
public static final String HANDLE_TIMESTAMP_WITHOUT_TIMEZONE =
"handle-timestamp-without-timezone";

// Output partition spec ID where writes should go.
public static final String OUTPUT_SPEC_ID = "output-spec-id";

public static final String OVERWRITE_MODE = "overwrite-mode";

// Overrides the default distribution mode for a write operation
Expand Down
24 changes: 20 additions & 4 deletions spark/v3.2/spark/src/main/java/org/apache/iceberg/spark/source/SparkWrite.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,7 @@
import org.apache.iceberg.io.OutputFileFactory;
import org.apache.iceberg.io.PartitioningWriter;
import org.apache.iceberg.io.RollingDataWriter;
import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
import org.apache.iceberg.relocated.com.google.common.collect.ImmutableSet;
import org.apache.iceberg.relocated.com.google.common.collect.Lists;
import org.apache.iceberg.spark.CommitMetadata;
Expand Down Expand Up @@ -90,6 +91,7 @@ abstract class SparkWrite implements Write, RequiresDistributionAndOrdering {
private final String applicationId;
private final boolean wapEnabled;
private final String wapId;
private final int outputSpecId;
private final long targetFileSize;
private final Schema writeSchema;
private final StructType dsSchema;
Expand Down Expand Up @@ -125,6 +127,15 @@ abstract class SparkWrite implements Write, RequiresDistributionAndOrdering {
this.partitionedFanoutEnabled = writeConf.fanoutWriterEnabled();
this.requiredDistribution = requiredDistribution;
this.requiredOrdering = requiredOrdering;

Preconditions.checkArgument(
writeConf.outputSpecId() == null || table.specs().containsKey(writeConf.outputSpecId()),
"Cannot write to unknown spec: %s",
writeConf.outputSpecId());
this.outputSpecId =
writeConf.outputSpecId() != null
? table.specs().get(writeConf.outputSpecId()).specId()
: table.spec().specId();
}

@Override
Expand Down Expand Up @@ -174,6 +185,7 @@ private WriterFactory createWriterFactory() {
tableBroadcast,
queryId,
format,
outputSpecId,
targetFileSize,
writeSchema,
dsSchema,
Expand Down Expand Up @@ -584,6 +596,7 @@ DataFile[] files() {
private static class WriterFactory implements DataWriterFactory, StreamingDataWriterFactory {
private final Broadcast<Table> tableBroadcast;
private final FileFormat format;
private final int outputSpecId;
private final long targetFileSize;
private final Schema writeSchema;
private final StructType dsSchema;
Expand All @@ -594,12 +607,14 @@ protected WriterFactory(
Broadcast<Table> tableBroadcast,
String queryId,
FileFormat format,
int outputSpecId,
long targetFileSize,
Schema writeSchema,
StructType dsSchema,
boolean partitionedFanoutEnabled) {
this.tableBroadcast = tableBroadcast;
this.format = format;
this.outputSpecId = outputSpecId;
this.targetFileSize = targetFileSize;
this.writeSchema = writeSchema;
this.dsSchema = dsSchema;
Expand All @@ -615,8 +630,8 @@ public DataWriter<InternalRow> createWriter(int partitionId, long taskId) {
@Override
public DataWriter<InternalRow> createWriter(int partitionId, long taskId, long epochId) {
Table table = tableBroadcast.value();
PartitionSpec spec = table.spec();
FileIO io = table.io();
PartitionSpec outputSpec = table.specs().get(outputSpecId);

OutputFileFactory fileFactory =
OutputFileFactory.builderFor(table, partitionId, taskId)
Expand All @@ -630,15 +645,16 @@ public DataWriter<InternalRow> createWriter(int partitionId, long taskId, long e
.dataSparkType(dsSchema)
.build();

if (spec.isUnpartitioned()) {
return new UnpartitionedDataWriter(writerFactory, fileFactory, io, spec, targetFileSize);
if (outputSpec.isUnpartitioned()) {
return new UnpartitionedDataWriter(
writerFactory, fileFactory, io, outputSpec, targetFileSize);

} else {
return new PartitionedDataWriter(
writerFactory,
fileFactory,
io,
spec,
outputSpec,
writeSchema,
dsSchema,
targetFileSize,
Expand Down
Loading