From 0ed49047213d2da5da563ec43deb501adc5e9fc9 Mon Sep 17 00:00:00 2001 From: Zihan Dai Date: Mon, 28 Sep 2026 19:33:56 +1000 Subject: [PATCH] [lance] Close the native and Arrow writers even when the final flush fails --- .../format/lance/LanceRecordsWriter.java | 55 +++++++++++-- .../format/lance/LanceRecordsWriterTest.java | 77 +++++++++++++++++++ 2 files changed, 124 insertions(+), 8 deletions(-) diff --git a/paimon-lance/src/main/java/org/apache/paimon/format/lance/LanceRecordsWriter.java b/paimon-lance/src/main/java/org/apache/paimon/format/lance/LanceRecordsWriter.java index a23d95489f9e..47b6f60ae8cf 100644 --- a/paimon-lance/src/main/java/org/apache/paimon/format/lance/LanceRecordsWriter.java +++ b/paimon-lance/src/main/java/org/apache/paimon/format/lance/LanceRecordsWriter.java @@ -95,9 +95,35 @@ public boolean reachTargetSize(boolean suggestedCheck, long targetSize) throws I @Override public void close() throws IOException { - flush(); + Throwable throwable = null; + + try { + flush(); + } catch (Throwable t) { + throwable = t; + } + LOG.info("Jni cost: " + jniCost + "ms for file: " + nativeWriter.path()); - closeImpl(); + long t1 = System.currentTimeMillis(); + + try { + nativeWriter.close(); + } catch (Throwable t) { + throwable = addSuppressed(throwable, t); + } + + try { + arrowFormatWriter.close(); + } catch (Throwable t) { + throwable = addSuppressed(throwable, t); + } + + long closeCost = (System.currentTimeMillis() - t1); + LOG.info("Close cost: " + closeCost + "ms for file: " + nativeWriter.path()); + + if (throwable != null) { + rethrow(throwable); + } } private void flush() throws IOException { @@ -110,11 +136,24 @@ private void flush() throws IOException { arrowFormatWriter.reset(); } - private void closeImpl() throws IOException { - long t1 = System.currentTimeMillis(); - this.nativeWriter.close(); - this.arrowFormatWriter.close(); - long closeCost = (System.currentTimeMillis() - t1); - LOG.info("Close cost: " + closeCost + "ms for file: " + nativeWriter.path()); + private static Throwable addSuppressed(Throwable throwable, Throwable suppressed) { + if (throwable == null) { + return suppressed; + } + throwable.addSuppressed(suppressed); + return throwable; + } + + private static void rethrow(Throwable throwable) throws IOException { + if (throwable instanceof IOException) { + throw (IOException) throwable; + } + if (throwable instanceof RuntimeException) { + throw (RuntimeException) throwable; + } + if (throwable instanceof Error) { + throw (Error) throwable; + } + throw new IOException(throwable); } } diff --git a/paimon-lance/src/test/java/org/apache/paimon/format/lance/LanceRecordsWriterTest.java b/paimon-lance/src/test/java/org/apache/paimon/format/lance/LanceRecordsWriterTest.java index 7d5161d5dca2..c1b8b5f9c89c 100644 --- a/paimon-lance/src/test/java/org/apache/paimon/format/lance/LanceRecordsWriterTest.java +++ b/paimon-lance/src/test/java/org/apache/paimon/format/lance/LanceRecordsWriterTest.java @@ -44,6 +44,7 @@ import java.util.stream.Collectors; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; /** Tests for {@link LanceRecordsWriter}. */ class LanceRecordsWriterTest { @@ -209,12 +210,88 @@ void testDifferentAllocatorRootFallsBackToRows() throws Exception { assertThat(nativeWriter.snapshots.get(0).values.get(0)).containsExactly(10); } + @Test + void testCloseContinuesCleanupAfterFlushFailure() { + RowType rowType = RowType.builder().field("value", DataTypes.INT()).build(); + RuntimeException flushFailure = new RuntimeException("flush failed"); + IOException nativeCloseFailure = new IOException("native close failed"); + IllegalStateException arrowCloseFailure = new IllegalStateException("arrow close failed"); + FailingCloseArrowFormatWriter arrowWriter = + new FailingCloseArrowFormatWriter(rowType, flushFailure, arrowCloseFailure); + FailingCloseLanceWriter nativeWriter = new FailingCloseLanceWriter(nativeCloseFailure); + LanceRecordsWriter writer = new LanceRecordsWriter(() -> 0L, arrowWriter, nativeWriter); + + try { + assertThatThrownBy(writer::close).isSameAs(flushFailure); + assertThat(flushFailure.getSuppressed()) + .containsExactly(nativeCloseFailure, arrowCloseFailure); + assertThat(nativeWriter.closeCount).isEqualTo(1); + assertThat(arrowWriter.closeCount).isEqualTo(1); + } finally { + if (arrowWriter.closeCount == 0) { + arrowWriter.closeWithoutFailure(); + } + } + } + private static void setInt(IntVector vector, int value) { vector.allocateNew(1); vector.setSafe(0, value); vector.setValueCount(1); } + private static class FailingCloseArrowFormatWriter extends ArrowFormatWriter { + + private final RuntimeException flushFailure; + private final RuntimeException closeFailure; + private int closeCount; + + private FailingCloseArrowFormatWriter( + RowType rowType, RuntimeException flushFailure, RuntimeException closeFailure) { + super(rowType, 1024, true); + this.flushFailure = flushFailure; + this.closeFailure = closeFailure; + } + + @Override + public void flush() { + throw flushFailure; + } + + @Override + public void close() { + closeCount++; + super.close(); + throw closeFailure; + } + + private void closeWithoutFailure() { + super.close(); + } + } + + private static class FailingCloseLanceWriter extends LanceWriter { + + private final IOException closeFailure; + private int closeCount; + + private FailingCloseLanceWriter(IOException closeFailure) { + super("unused", Collections.emptyMap()); + this.closeFailure = closeFailure; + } + + @Override + public void close() throws IOException { + closeCount++; + throw closeFailure; + } + + @Override + public String path() { + return "unused"; + } + } + private static class CapturingLanceWriter extends LanceWriter { private final List snapshots = new ArrayList<>();