[lance] Close the native and Arrow writers even when the final flush fails

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);
     }
 }

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<Snapshot> snapshots = new ArrayList<>();