Skip to content

Commit ec6c7fb

Browse files
gitRasheedgitRasheed
authored
fix(spark): let abort close every writer after a failed write (#11084)
When a write task fails, Spark's concurrent partition writer closes the task's open writers one after another and stops at the first `close()` that throws. `VortexOutputWriter.close()` retried the batch that had just failed, so it failed again and threw before releasing anything, the native writer included. Spark then never closed the remaining writers, and they kept their Arrow memory, output streams and native writers. --------- Signed-off-by: gitRasheed <[email protected]> Signed-off-by: Robert Kruszewski <[email protected]>
1 parent 75a4b92 commit ec6c7fb

4 files changed

Lines changed: 243 additions & 7 deletions

File tree

‎java/buildSrc/src/main/kotlin/vortex-spark-module.gradle.kts‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -159,6 +159,12 @@ testing {
159159
versionSourceDir.resolve("test/java"),
160160
),
161161
)
162+
scala.setSrcDirs(
163+
listOf(
164+
commonSourceDir.resolve("test/scala"),
165+
versionSourceDir.resolve("test/scala"),
166+
),
167+
)
162168
resources.setSrcDirs(
163169
listOf(
164170
commonSourceDir.resolve("test/resources"),
@@ -243,6 +249,8 @@ val sparkJvmArgs =
243249
"--add-opens=java.base/java.nio=ALL-UNNAMED",
244250
"--add-opens=java.base/sun.util.calendar=ALL-UNNAMED",
245251
"--add-opens=java.base/sun.security.action=ALL-UNNAMED",
252+
// Spark 4.2's SerializationDebugger fails to initialize without it when a task failure won't serialize.
253+
"--add-opens=java.base/java.io=ALL-UNNAMED",
246254
)
247255

248256
tasks.withType<Test>().configureEach {

‎java/vortex-spark/common/src/main/java/dev/vortex/spark/write/VortexOutputWriter.java‎

Lines changed: 53 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535
import java.io.UncheckedIOException;
3636
import java.util.ArrayList;
3737
import java.util.List;
38+
import org.apache.spark.TaskContext;
3839
import org.apache.spark.sql.catalyst.InternalRow;
3940
import org.apache.spark.sql.catalyst.expressions.SpecializedGetters;
4041
import org.apache.spark.sql.catalyst.util.ArrayData;
@@ -92,6 +93,7 @@ public final class VortexOutputWriter extends OutputWriter {
9293
private BufferAllocator allocator;
9394
private VectorSchemaRoot vectorSchemaRoot;
9495
private final List<InternalRow> batchRows = new ArrayList<>();
96+
private Exception writeFailure;
9597
private boolean closed = false;
9698

9799
/**
@@ -165,6 +167,10 @@ static int configuredBatchSize(VortexOptions options) {
165167
*/
166168
@Override
167169
public void write(InternalRow row) {
170+
if (writeFailure != null) {
171+
throw new IllegalStateException(
172+
"VortexOutputWriter for " + filePath + " failed on an earlier write", writeFailure);
173+
}
168174
// Add row to current batch
169175
batchRows.add(row.copy());
170176

@@ -173,7 +179,11 @@ public void write(InternalRow row) {
173179
try {
174180
writeBatch();
175181
} catch (IOException e) {
182+
writeFailure = e;
176183
throw new UncheckedIOException("Failed to write a Vortex batch to " + filePath, e);
184+
} catch (RuntimeException e) {
185+
writeFailure = e;
186+
throw e;
177187
}
178188
}
179189
}
@@ -326,26 +336,44 @@ private void populateStructVector(StructVector vector, StructType dataType, Inte
326336
}
327337
}
328338

339+
/**
340+
* Flushes the last batch and finalizes the file.
341+
*
342+
* <p>After a failed {@link #write}, or in a task that has already failed, the last batch is dropped because the
343+
* file is never committed. In both cases this releases everything and logs what fails instead of throwing, because
344+
* Spark's abort closes a task's writers one after another and stops at the first that throws.
345+
*/
329346
@Override
330347
public void close() {
331348
if (closed) {
332349
return;
333350
}
351+
352+
boolean aborting = writeFailure != null || taskFailed();
334353
IOException failure = null;
335-
try {
336-
if (!batchRows.isEmpty()) {
354+
355+
if (!aborting && !batchRows.isEmpty()) {
356+
try {
337357
writeBatch();
358+
} catch (IOException e) {
359+
failure = e;
360+
} catch (RuntimeException e) {
361+
failure = new IOException("Failed to write the last Vortex batch", e);
338362
}
339-
if (vortexWriter != null) {
363+
}
364+
365+
if (vortexWriter != null) {
366+
try {
340367
vortexWriter.finish();
368+
} catch (IOException e) {
369+
failure = addFailure(failure, e);
370+
} finally {
371+
vortexWriter = null;
341372
}
342-
} catch (IOException e) {
343-
failure = e;
344-
} finally {
345-
vortexWriter = null;
346373
}
347374

348375
failure = closeArrowResources(failure);
376+
349377
if (writable != null) {
350378
try {
351379
writable.close();
@@ -355,8 +383,15 @@ public void close() {
355383
writable = null;
356384
}
357385
}
386+
358387
session = null;
359388
closed = true;
389+
390+
if (failure != null && aborting) {
391+
logger.warn("Failed to close the writer for {} in a failed task", filePath, failure);
392+
return;
393+
}
394+
360395
if (failure != null) {
361396
throw new UncheckedIOException("Failed to close Vortex output " + filePath, failure);
362397
}
@@ -367,6 +402,17 @@ public String path() {
367402
return filePath;
368403
}
369404

405+
/** Spark marks a task failed before its abort closes the task's writers. */
406+
private static boolean taskFailed() {
407+
TaskContext context = TaskContext.get();
408+
return context != null && context.isFailed();
409+
}
410+
411+
/** Package-private so tests can see what close leaves allocated. */
412+
BufferAllocator allocator() {
413+
return allocator;
414+
}
415+
370416
private IOException closeArrowResources(IOException failure) {
371417
if (vectorSchemaRoot != null) {
372418
try {

‎java/vortex-spark/common/src/test/java/dev/vortex/spark/write/VortexOutputWriterAbortTest.java‎

Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3+
4+
package dev.vortex.spark.write;
5+
6+
import static org.junit.jupiter.api.Assertions.assertEquals;
7+
import static org.junit.jupiter.api.Assertions.assertThrows;
8+
9+
import dev.vortex.relocated.org.apache.arrow.memory.BufferAllocator;
10+
import java.io.IOException;
11+
import java.io.PrintWriter;
12+
import java.io.StringWriter;
13+
import java.net.URI;
14+
import java.nio.file.Files;
15+
import java.nio.file.Path;
16+
import java.util.ArrayList;
17+
import java.util.List;
18+
import java.util.stream.Stream;
19+
import org.apache.spark.sql.Row;
20+
import org.apache.spark.sql.RowFactory;
21+
import org.apache.spark.sql.SparkSession;
22+
import org.apache.spark.sql.types.DataTypes;
23+
import org.apache.spark.sql.types.StructType;
24+
import org.junit.jupiter.api.io.TempDir;
25+
import org.junit.jupiter.params.ParameterizedTest;
26+
import org.junit.jupiter.params.provider.CsvSource;
27+
28+
/**
29+
* Spark's concurrent partition writer closes its open writers one after another when a task aborts and stops at the
30+
* first close that throws, so a writer whose write already failed must not throw again from close.
31+
*/
32+
final class VortexOutputWriterAbortTest {
33+
private static final List<VortexOutputWriter> WRITERS = new ArrayList<>();
34+
private static final List<BufferAllocator> ALLOCATORS = new ArrayList<>();
35+
private static final int PARTITIONS = 8;
36+
37+
@TempDir
38+
Path tempDir;
39+
40+
// With one-row batches an oversized row fails in write, otherwise in the close that flushes it. With one in every
41+
// partition, the closes during abort fail too.
42+
@ParameterizedTest
43+
@CsvSource({"1, false", "2048, false", "2048, true"})
44+
void aFailedWriteLetsAbortCloseEveryPartitionWriter(String batchSize, boolean everyPartition) throws IOException {
45+
SparkSession spark = SparkSession.builder()
46+
.master("local[1,1]")
47+
.appName("VortexOutputWriterAbortTest")
48+
.config("spark.driver.host", "127.0.0.1")
49+
.config("spark.ui.enabled", "false")
50+
.getOrCreate();
51+
spark.conf().set("spark.sql.adaptive.enabled", "false");
52+
spark.conf().set("spark.sql.maxConcurrentOutputFileWriters", "16");
53+
StructType schema = new StructType().add("s", DataTypes.StringType).add("p", DataTypes.IntegerType);
54+
List<String> stranded = new ArrayList<>();
55+
try {
56+
// The order Spark closes its writers in decides whether the failed one comes first, so fail each in turn.
57+
for (int failed = 0; failed < (everyPartition ? 1 : PARTITIONS); failed++) {
58+
WRITERS.clear();
59+
ALLOCATORS.clear();
60+
List<Row> rows = new ArrayList<>();
61+
for (int p = 0; p < PARTITIONS; p++) {
62+
rows.add(RowFactory.create("small", p));
63+
}
64+
// Larger than the allocator limit set below, so appending it fails.
65+
for (int p = 0; p < PARTITIONS; p++) {
66+
if (everyPartition || p == failed) {
67+
rows.add(RowFactory.create("x".repeat(2 << 20), p));
68+
}
69+
}
70+
Path output = tempDir.resolve(batchSize + "-" + everyPartition + "-failed=" + failed);
71+
72+
Exception thrown = assertThrows(Exception.class, () -> spark.createDataFrame(rows, schema)
73+
.coalesce(1)
74+
.write()
75+
.format("dev.vortex.spark.write.LimitedVortexFileFormat")
76+
.option("vortex.write.batch.size", batchSize)
77+
.partitionBy("p")
78+
.save(output.toString()));
79+
80+
assertEquals(PARTITIONS, WRITERS.size(), () -> stackTrace(thrown));
81+
for (int i = 0; i < WRITERS.size(); i++) {
82+
VortexOutputWriter writer = WRITERS.get(i);
83+
long held = ALLOCATORS.get(i).getAllocatedMemory();
84+
if (writer.allocator() != null || held != 0) {
85+
String partition = Path.of(URI.create(writer.path()))
86+
.getParent()
87+
.getFileName()
88+
.toString();
89+
stranded.add("failed=" + failed + " " + partition + ": " + held + " bytes");
90+
try {
91+
writer.close();
92+
} catch (RuntimeException ignored) {
93+
// Only the release matters here; the assertion below reports the stranding.
94+
}
95+
}
96+
}
97+
if (Files.exists(output)) {
98+
try (Stream<Path> files = Files.walk(output)) {
99+
assertEquals(
100+
List.of(),
101+
files.filter(f -> f.toString().endsWith(".vortex") || f.endsWith("_SUCCESS"))
102+
.toList());
103+
}
104+
}
105+
}
106+
} finally {
107+
spark.stop();
108+
WRITERS.clear();
109+
ALLOCATORS.clear();
110+
}
111+
assertEquals(List.of(), stranded, "writers left open after the task aborted");
112+
}
113+
114+
private static String stackTrace(Throwable t) {
115+
StringWriter out = new StringWriter();
116+
t.printStackTrace(new PrintWriter(out));
117+
return out.toString();
118+
}
119+
120+
/** Called for every writer LimitedVortexFileFormat opens, which is Scala and compiled after this class. */
121+
static VortexOutputWriter opened(VortexOutputWriter writer) {
122+
writer.allocator().setLimit(1 << 20);
123+
WRITERS.add(writer);
124+
ALLOCATORS.add(writer.allocator());
125+
return writer;
126+
}
127+
}

‎java/vortex-spark/common/src/test/scala/dev/vortex/spark/write/LimitedVortexFileFormat.scala‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3+
4+
package dev.vortex.spark.write
5+
6+
import org.apache.hadoop.fs.FileStatus
7+
import org.apache.hadoop.mapreduce.{Job, TaskAttemptContext}
8+
9+
import org.apache.spark.sql.SparkSession
10+
import org.apache.spark.sql.execution.datasources.{
11+
FileFormat,
12+
OutputWriter,
13+
OutputWriterFactory
14+
}
15+
import org.apache.spark.sql.types.StructType
16+
17+
import dev.vortex.spark.VortexFileFormat
18+
19+
/** Vortex, with every writer handed to [[VortexOutputWriterAbortTest]]. In
20+
* Scala because FileFormat's Seq differs between Scala 2.12 and 2.13.
21+
*/
22+
final class LimitedVortexFileFormat extends FileFormat {
23+
private val delegate = new VortexFileFormat
24+
25+
override def inferSchema(
26+
spark: SparkSession,
27+
options: Map[String, String],
28+
files: Seq[FileStatus]
29+
): Option[StructType] = delegate.inferSchema(spark, options, files)
30+
31+
override def prepareWrite(
32+
spark: SparkSession,
33+
job: Job,
34+
options: Map[String, String],
35+
dataSchema: StructType
36+
): OutputWriterFactory = new LimitedOutputWriterFactory(
37+
delegate.prepareWrite(spark, job, options, dataSchema)
38+
)
39+
}
40+
41+
final class LimitedOutputWriterFactory(delegate: OutputWriterFactory)
42+
extends OutputWriterFactory {
43+
override def getFileExtension(context: TaskAttemptContext): String =
44+
delegate.getFileExtension(context)
45+
46+
override def newInstance(
47+
path: String,
48+
dataSchema: StructType,
49+
context: TaskAttemptContext
50+
): OutputWriter = VortexOutputWriterAbortTest.opened(
51+
delegate
52+
.newInstance(path, dataSchema, context)
53+
.asInstanceOf[VortexOutputWriter]
54+
)
55+
}

0 commit comments

Comments
 (0)