org.dinky.shaded.paimon.mergetree.compact.ChangelogMergeTreeRewriter Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of dinky-shaded-paimon Show documentation
The newest version!
/*
 * 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.dinky.shaded.paimon.mergetree.compact;

import org.dinky.shaded.paimon.KeyValue;
import org.dinky.shaded.paimon.codegen.RecordEqualiser;
import org.dinky.shaded.paimon.compact.CompactResult;
import org.dinky.shaded.paimon.data.InternalRow;
import org.dinky.shaded.paimon.io.DataFileMeta;
import org.dinky.shaded.paimon.io.KeyValueFileReaderFactory;
import org.dinky.shaded.paimon.io.KeyValueFileWriterFactory;
import org.dinky.shaded.paimon.io.RollingFileWriter;
import org.dinky.shaded.paimon.mergetree.MergeSorter;
import org.dinky.shaded.paimon.mergetree.MergeTreeReaders;
import org.dinky.shaded.paimon.mergetree.SortedRun;
import org.dinky.shaded.paimon.reader.RecordReaderIterator;

import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;

/** A {@link MergeTreeCompactRewriter} which produces changelog files for the compaction. */
public abstract class ChangelogMergeTreeRewriter extends MergeTreeCompactRewriter {

    protected final RecordEqualiser valueEqualiser;
    protected final boolean changelogRowDeduplicate;

    public ChangelogMergeTreeRewriter(
            KeyValueFileReaderFactory readerFactory,
            KeyValueFileWriterFactory writerFactory,
            Comparator keyComparator,
            MergeFunctionFactory mfFactory,
            MergeSorter mergeSorter,
            RecordEqualiser valueEqualiser,
            boolean changelogRowDeduplicate) {
        super(readerFactory, writerFactory, keyComparator, mfFactory, mergeSorter);
        this.valueEqualiser = valueEqualiser;
        this.changelogRowDeduplicate = changelogRowDeduplicate;
    }

    protected abstract boolean rewriteChangelog(
            int outputLevel, boolean dropDelete, List> sections);

    protected abstract boolean upgradeChangelog(int outputLevel, DataFileMeta file);

    protected abstract MergeFunctionWrapper createMergeWrapper(int outputLevel);

    protected boolean rewriteLookupChangelog(int outputLevel, List> sections) {
        if (outputLevel == 0) {
            return false;
        }

        for (List runs : sections) {
            for (SortedRun run : runs) {
                for (DataFileMeta file : run.files()) {
                    if (file.level() == 0) {
                        return true;
                    }
                }
            }
        }
        return false;
    }

    @Override
    public CompactResult rewrite(
            int outputLevel, boolean dropDelete, List> sections) throws Exception {
        if (rewriteChangelog(outputLevel, dropDelete, sections)) {
            return rewriteChangelogCompaction(outputLevel, sections);
        } else {
            return rewriteCompaction(outputLevel, dropDelete, sections);
        }
    }

    private CompactResult rewriteChangelogCompaction(
            int outputLevel, List> sections) throws Exception {
        List> sectionReaders = new ArrayList<>();
        for (List section : sections) {
            sectionReaders.add(
                    () ->
                            MergeTreeReaders.readerForSection(
                                    section,
                                    readerFactory,
                                    keyComparator,
                                    createMergeWrapper(outputLevel),
                                    mergeSorter));
        }

        RecordReaderIterator iterator = null;
        RollingFileWriter compactFileWriter = null;
        RollingFileWriter changelogFileWriter = null;

        try {
            iterator = new RecordReaderIterator<>(ConcatRecordReader.create(sectionReaders));
            compactFileWriter = writerFactory.createRollingMergeTreeFileWriter(outputLevel);
            changelogFileWriter = writerFactory.createRollingChangelogFileWriter(outputLevel);

            while (iterator.hasNext()) {
                ChangelogResult result = iterator.next();
                if (result.result() != null) {
                    compactFileWriter.write(result.result());
                }
                for (KeyValue kv : result.changelogs()) {
                    changelogFileWriter.write(kv);
                }
            }
        } finally {
            if (iterator != null) {
                iterator.close();
            }
            if (compactFileWriter != null) {
                compactFileWriter.close();
            }
            if (changelogFileWriter != null) {
                changelogFileWriter.close();
            }
        }

        return new CompactResult(
                extractFilesFromSections(sections),
                compactFileWriter.result(),
                changelogFileWriter.result());
    }

    @Override
    public CompactResult upgrade(int outputLevel, DataFileMeta file) throws Exception {
        if (upgradeChangelog(outputLevel, file)) {
            return rewriteChangelogCompaction(
                    outputLevel,
                    Collections.singletonList(
                            Collections.singletonList(SortedRun.fromSingle(file))));
        } else {
            return super.upgrade(outputLevel, file);
        }
    }
}