From 097d3e7f9afa272cd59a2805ee044ac7257ada5a Mon Sep 17 00:00:00 2001 From: umi Date: Tue, 3 Nov 2026 16:52:54 +0800 Subject: [PATCH] [core] Release lookup merge spills when closing sort readers Reset the merge wrapper when either sort reader closes so the last key releases its spill files. Preserve input close failures when resetting the wrapper also fails. --- .../compact/SortMergeReaderWithLoserTree.java | 8 +- .../compact/SortMergeReaderWithMinHeap.java | 3 +- .../compact/SortMergeReaderCleanupTest.java | 154 ++++++++++++++++++ 3 files changed, 163 insertions(+), 2 deletions(-) create mode 100644 paimon-core/src/test/java/org/apache/paimon/mergetree/compact/SortMergeReaderCleanupTest.java 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(); + } + } +}