diff --git a/paimon-core/src/main/java/org/apache/paimon/mergetree/compact/SortMergeReaderWithLoserTree.java b/paimon-core/src/main/java/org/apache/paimon/mergetree/compact/SortMergeReaderWithLoserTree.java index a3b8b1e26f59..256ca4f3835b 100644 --- a/paimon-core/src/main/java/org/apache/paimon/mergetree/compact/SortMergeReaderWithLoserTree.java +++ b/paimon-core/src/main/java/org/apache/paimon/mergetree/compact/SortMergeReaderWithLoserTree.java @@ -21,7 +21,9 @@ import org.apache.paimon.KeyValue; import org.apache.paimon.data.InternalRow; import org.apache.paimon.reader.RecordReader; +import org.apache.paimon.utils.ExceptionUtils; import org.apache.paimon.utils.FieldsComparator; +import org.apache.paimon.utils.IOUtils; import org.apache.paimon.utils.Preconditions; import javax.annotation.Nullable; @@ -84,7 +86,11 @@ public RecordIterator readBatch() throws IOException { @Override public void close() throws IOException { - loserTree.close(); + try { + IOUtils.closeAll(loserTree, mergeFunctionWrapper::reset); + } catch (Exception e) { + ExceptionUtils.rethrowIOException(e); + } } /** The iterator iterates on {@link SortMergeReaderWithLoserTree}. */ diff --git a/paimon-core/src/main/java/org/apache/paimon/mergetree/compact/SortMergeReaderWithMinHeap.java b/paimon-core/src/main/java/org/apache/paimon/mergetree/compact/SortMergeReaderWithMinHeap.java index db09ce41a7dd..ddba5897dd92 100644 --- a/paimon-core/src/main/java/org/apache/paimon/mergetree/compact/SortMergeReaderWithMinHeap.java +++ b/paimon-core/src/main/java/org/apache/paimon/mergetree/compact/SortMergeReaderWithMinHeap.java @@ -119,6 +119,7 @@ public void close() throws IOException { for (Element element : polled) { collected = releaseCollecting(element, collected); } + collected = closeCollecting(mergeFunctionWrapper::reset, collected); if (collected != null) { rethrowCloseFailure(collected); } @@ -135,7 +136,7 @@ public void close() throws IOException { } private static @Nullable Throwable closeCollecting( - RecordReader reader, @Nullable Throwable collected) { + AutoCloseable reader, @Nullable Throwable collected) { try { reader.close(); } catch (Throwable t) { diff --git a/paimon-core/src/test/java/org/apache/paimon/mergetree/compact/SortMergeReaderCleanupTest.java b/paimon-core/src/test/java/org/apache/paimon/mergetree/compact/SortMergeReaderCleanupTest.java new file mode 100644 index 000000000000..c045f2537581 --- /dev/null +++ b/paimon-core/src/test/java/org/apache/paimon/mergetree/compact/SortMergeReaderCleanupTest.java @@ -0,0 +1,154 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.paimon.mergetree.compact; + +import org.apache.paimon.CoreOptions; +import org.apache.paimon.CoreOptions.SortEngine; +import org.apache.paimon.KeyValue; +import org.apache.paimon.data.BinaryString; +import org.apache.paimon.data.GenericRow; +import org.apache.paimon.disk.IOManager; +import org.apache.paimon.options.MemorySize; +import org.apache.paimon.options.Options; +import org.apache.paimon.reader.RecordReader; +import org.apache.paimon.types.DataTypes; +import org.apache.paimon.types.RowKind; +import org.apache.paimon.types.RowType; +import org.apache.paimon.utils.IteratorRecordReader; + +import org.junit.jupiter.api.io.TempDir; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.EnumSource; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.Comparator; +import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.stream.Stream; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** Tests for merge state cleanup when a {@link SortMergeReader} is closed. */ +class SortMergeReaderCleanupTest { + + @TempDir Path tempDir; + + @ParameterizedTest + @EnumSource(SortEngine.class) + void testCloseReleasesLookupSpills(SortEngine sortEngine) throws Exception { + Options rawOptions = new Options(); + rawOptions.set(CoreOptions.LOOKUP_MERGE_RECORDS_THRESHOLD, 0); + rawOptions.set(CoreOptions.LOOKUP_MERGE_BUFFER_SIZE, MemorySize.parse("4 kb")); + rawOptions.set(CoreOptions.PAGE_SIZE, MemorySize.parse("4 kb")); + CoreOptions options = new CoreOptions(rawOptions); + RowType keyType = RowType.builder().field("k", DataTypes.INT()).build(); + RowType valueType = RowType.builder().field("v", DataTypes.STRING()).build(); + char[] chars = new char[1000]; + Arrays.fill(chars, 'x'); + BinaryString value = BinaryString.fromString(new String(chars)); + try (IOManager ioManager = IOManager.create(tempDir.toString())) { + MergeFunction merge = + new LookupMergeFunction( + DeduplicateMergeFunction.factory().create(), + options, + keyType, + valueType, + ioManager); + List> inputs = new ArrayList<>(); + for (int i = 0; i < 16; i++) { + KeyValue kv = + new KeyValue() + .replace(GenericRow.of(1), i, RowKind.INSERT, GenericRow.of(value)) + .setLevel(0); + inputs.add(new IteratorRecordReader<>(Collections.singletonList(kv).iterator())); + } + KeyValue result; + try (SortMergeReader reader = + SortMergeReader.createSortMergeReader( + inputs, + Comparator.comparingInt(row -> row.getInt(0)), + null, + new ReducerMergeFunctionWrapper(merge), + sortEngine)) { + RecordReader.RecordIterator batch = reader.readBatch(); + assertThat(batch).isNotNull(); + try { + result = batch.next(); + assertThat(result).isNotNull(); + assertThat(result.sequenceNumber()).isEqualTo(15); + assertThat(spillFileCount()).isPositive(); + } finally { + batch.releaseBatch(); + } + } + assertThat(spillFileCount()).isZero(); + assertThat(result.value().getString(0)).isEqualTo(value); + } + } + + @ParameterizedTest + @EnumSource(SortEngine.class) + void testCloseStillResetsMergeWhenInputCloseFails(SortEngine sortEngine) { + IOException inputFailure = new IOException("input close failed"); + RuntimeException resetFailure = new IllegalStateException("merge reset failed"); + RecordReader input = + new RecordReader() { + @Override + public RecordIterator readBatch() { + return null; + } + + @Override + public void close() throws IOException { + throw inputFailure; + } + }; + AtomicInteger resets = new AtomicInteger(); + MergeFunctionWrapper merge = + new ReducerMergeFunctionWrapper(DeduplicateMergeFunction.factory().create()) { + @Override + public void reset() { + resets.incrementAndGet(); + throw resetFailure; + } + }; + SortMergeReader reader = + SortMergeReader.createSortMergeReader( + Collections.singletonList(input), + Comparator.comparingInt(row -> row.getInt(0)), + null, + merge, + sortEngine); + assertThatThrownBy(reader::close).isSameAs(inputFailure); + assertThat(inputFailure.getSuppressed()).containsExactly(resetFailure); + assertThat(resets.get()).isEqualTo(1); + } + + private long spillFileCount() throws IOException { + try (Stream paths = Files.walk(tempDir)) { + return paths.filter(Files::isRegularFile).count(); + } + } +}