org.dinky.shaded.paimon.mergetree.LookupLevels Maven / Gradle / Ivy

Go to download
/*
 * 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;

import org.dinky.shaded.paimon.KeyValue;
import org.dinky.shaded.paimon.annotation.VisibleForTesting;
import org.dinky.shaded.paimon.data.InternalRow;
import org.dinky.shaded.paimon.data.serializer.RowCompactedSerializer;
import org.dinky.shaded.paimon.io.DataFileMeta;
import org.dinky.shaded.paimon.io.DataOutputSerializer;
import org.dinky.shaded.paimon.lookup.LookupStoreFactory;
import org.dinky.shaded.paimon.lookup.LookupStoreReader;
import org.dinky.shaded.paimon.lookup.LookupStoreWriter;
import org.dinky.shaded.paimon.memory.MemorySegment;
import org.dinky.shaded.paimon.options.MemorySize;
import org.dinky.shaded.paimon.reader.RecordReader;
import org.dinky.shaded.paimon.types.RowKind;
import org.dinky.shaded.paimon.types.RowType;
import org.dinky.shaded.paimon.utils.FileIOUtils;
import org.dinky.shaded.paimon.utils.IOFunction;

import org.dinky.shaded.paimon.shade.caffeine2.com.github.benmanes.caffeine.cache.Cache;
import org.dinky.shaded.paimon.shade.caffeine2.com.github.benmanes.caffeine.cache.Caffeine;
import org.dinky.shaded.paimon.shade.caffeine2.com.github.benmanes.caffeine.cache.RemovalCause;
import org.dinky.shaded.paimon.shade.guava30.com.google.common.util.concurrent.MoreExecutors;

import javax.annotation.Nullable;

import java.io.Closeable;
import java.io.File;
import java.io.IOException;
import java.io.UncheckedIOException;
import java.time.Duration;
import java.util.Comparator;
import java.util.function.Supplier;

import static org.dinky.shaded.paimon.mergetree.LookupUtils.fileKibiBytes;
import static org.dinky.shaded.paimon.utils.Preconditions.checkArgument;

/** Provide lookup by key. */
public class LookupLevels implements Levels.DropFileCallback, Closeable {

    private final Levels levels;
    private final Comparator keyComparator;
    private final RowCompactedSerializer keySerializer;
    private final RowCompactedSerializer valueSerializer;
    private final IOFunction> fileReaderFactory;
    private final Supplier localFileFactory;
    private final LookupStoreFactory lookupStoreFactory;

    private final Cache lookupFiles;

    public LookupLevels(
            Levels levels,
            Comparator keyComparator,
            RowType keyType,
            RowType valueType,
            IOFunction> fileReaderFactory,
            Supplier localFileFactory,
            LookupStoreFactory lookupStoreFactory,
            Duration fileRetention,
            MemorySize maxDiskSize) {
        this.levels = levels;
        this.keyComparator = keyComparator;
        this.keySerializer = new RowCompactedSerializer(keyType);
        this.valueSerializer = new RowCompactedSerializer(valueType);
        this.fileReaderFactory = fileReaderFactory;
        this.localFileFactory = localFileFactory;
        this.lookupStoreFactory = lookupStoreFactory;
        this.lookupFiles =
                Caffeine.newBuilder()
                        .expireAfterAccess(fileRetention)
                        .maximumWeight(maxDiskSize.getKibiBytes())
                        .weigher(this::fileWeigh)
                        .removalListener(this::removalCallback)
                        .executor(MoreExecutors.directExecutor())
                        .build();
        levels.addDropFileCallback(this);
    }

    @VisibleForTesting
    Cache lookupFiles() {
        return lookupFiles;
    }

    @Override
    public void notifyDropFile(String file) {
        lookupFiles.invalidate(file);
    }

    @Nullable
    public KeyValue lookup(InternalRow key, int startLevel) throws IOException {
        return LookupUtils.lookup(levels, key, startLevel, this::lookup);
    }

    @Nullable
    private KeyValue lookup(InternalRow key, SortedRun level) throws IOException {
        return LookupUtils.lookup(keyComparator, key, level, this::lookup);
    }

    @Nullable
    private KeyValue lookup(InternalRow key, DataFileMeta file) throws IOException {
        LookupFile lookupFile = lookupFiles.getIfPresent(file.fileName());
        while (lookupFile == null || lookupFile.isClosed) {
            lookupFile = createLookupFile(file);
            lookupFiles.put(file.fileName(), lookupFile);
        }

        byte[] keyBytes = keySerializer.serializeToBytes(key);
        byte[] valueBytes = lookupFile.get(keyBytes);
        if (valueBytes == null) {
            return null;
        }
        InternalRow value = valueSerializer.deserialize(valueBytes);
        long sequenceNumber = MemorySegment.wrap(valueBytes).getLong(valueBytes.length - 9);
        RowKind rowKind = RowKind.fromByteValue(valueBytes[valueBytes.length - 1]);
        return new KeyValue()
                .replace(key, sequenceNumber, rowKind, value)
                .setLevel(lookupFile.remoteFile().level());
    }

    private int fileWeigh(String file, LookupFile lookupFile) {
        return fileKibiBytes(lookupFile.localFile);
    }

    private void removalCallback(String key, LookupFile file, RemovalCause cause) {
        if (file != null) {
            try {
                file.close();
            } catch (IOException e) {
                throw new UncheckedIOException(e);
            }
        }
    }

    private LookupFile createLookupFile(DataFileMeta file) throws IOException {
        File localFile = localFileFactory.get();
        if (!localFile.createNewFile()) {
            throw new IOException("Can not create new file: " + localFile);
        }
        try (LookupStoreWriter kvWriter = lookupStoreFactory.createWriter(localFile);
                RecordReader reader = fileReaderFactory.apply(file)) {
            DataOutputSerializer valueOut = new DataOutputSerializer(32);
            RecordReader.RecordIterator batch;
            KeyValue kv;
            while ((batch = reader.readBatch()) != null) {
                while ((kv = batch.next()) != null) {
                    byte[] keyBytes = keySerializer.serializeToBytes(kv.key());
                    valueOut.clear();
                    valueOut.write(valueSerializer.serializeToBytes(kv.value()));
                    valueOut.writeLong(kv.sequenceNumber());
                    valueOut.writeByte(kv.valueKind().toByteValue());
                    byte[] valueBytes = valueOut.getCopyOfBuffer();
                    kvWriter.put(keyBytes, valueBytes);
                }
                batch.releaseBatch();
            }
        } catch (IOException e) {
            FileIOUtils.deleteFileOrDirectory(localFile);
            throw e;
        }

        return new LookupFile(localFile, file, lookupStoreFactory.createReader(localFile));
    }

    @Override
    public void close() throws IOException {
        lookupFiles.invalidateAll();
    }

    private static class LookupFile implements Closeable {

        private final File localFile;
        private final DataFileMeta remoteFile;
        private final LookupStoreReader reader;

        private boolean isClosed = false;

        public LookupFile(File localFile, DataFileMeta remoteFile, LookupStoreReader reader) {
            this.localFile = localFile;
            this.remoteFile = remoteFile;
            this.reader = reader;
        }

        @Nullable
        public byte[] get(byte[] key) throws IOException {
            checkArgument(!isClosed);
            return reader.lookup(key);
        }

        public DataFileMeta remoteFile() {
            return remoteFile;
        }

        @Override
        public void close() throws IOException {
            reader.close();
            isClosed = true;
            FileIOUtils.deleteFileOrDirectory(localFile);
        }
    }
}