org.dinky.shaded.paimon.mergetree.ContainsLevels Maven / Gradle / Ivy
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;
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.lookup.LookupStoreFactory;
import org.dinky.shaded.paimon.lookup.LookupStoreReader;
import org.dinky.shaded.paimon.lookup.LookupStoreWriter;
import org.dinky.shaded.paimon.options.MemorySize;
import org.dinky.shaded.paimon.reader.RecordReader;
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 contains key. */
public class ContainsLevels implements Levels.DropFileCallback, Closeable {
private static final byte[] EMPTY_VALUE = new byte[0];
private final Levels levels;
private final Comparator keyComparator;
private final RowCompactedSerializer keySerializer;
private final IOFunction> fileReaderFactory;
private final Supplier localFileFactory;
private final LookupStoreFactory lookupStoreFactory;
private final Cache containsFiles;
public ContainsLevels(
Levels levels,
Comparator keyComparator,
RowType keyType,
IOFunction> fileReaderFactory,
Supplier localFileFactory,
LookupStoreFactory lookupStoreFactory,
Duration fileRetention,
MemorySize maxDiskSize) {
this.levels = levels;
this.keyComparator = keyComparator;
this.keySerializer = new RowCompactedSerializer(keyType);
this.fileReaderFactory = fileReaderFactory;
this.localFileFactory = localFileFactory;
this.lookupStoreFactory = lookupStoreFactory;
this.containsFiles =
Caffeine.newBuilder()
.expireAfterAccess(fileRetention)
.maximumWeight(maxDiskSize.getKibiBytes())
.weigher(this::fileWeigh)
.removalListener(this::removalCallback)
.executor(MoreExecutors.directExecutor())
.build();
levels.addDropFileCallback(this);
}
@VisibleForTesting
Cache containsFiles() {
return containsFiles;
}
@Override
public void notifyDropFile(String file) {
containsFiles.invalidate(file);
}
public boolean contains(InternalRow key, int startLevel) throws IOException {
Boolean result = LookupUtils.lookup(levels, key, startLevel, this::contains);
return result != null && result;
}
@Nullable
private Boolean contains(InternalRow key, SortedRun level) throws IOException {
return LookupUtils.lookup(keyComparator, key, level, this::contains);
}
@Nullable
private Boolean contains(InternalRow key, DataFileMeta file) throws IOException {
ContainsFile containsFile = containsFiles.getIfPresent(file.fileName());
while (containsFile == null || containsFile.isClosed) {
containsFile = createContainsFile(file);
containsFiles.put(file.fileName(), containsFile);
}
if (containsFile.get(keySerializer.serializeToBytes(key)) != null) {
return true;
}
return null;
}
private int fileWeigh(String file, ContainsFile containsFile) {
return fileKibiBytes(containsFile.localFile);
}
private void removalCallback(String key, ContainsFile file, RemovalCause cause) {
if (file != null) {
try {
file.close();
} catch (IOException e) {
throw new UncheckedIOException(e);
}
}
}
private ContainsFile createContainsFile(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)) {
RecordReader.RecordIterator batch;
KeyValue kv;
while ((batch = reader.readBatch()) != null) {
while ((kv = batch.next()) != null) {
byte[] keyBytes = keySerializer.serializeToBytes(kv.key());
kvWriter.put(keyBytes, EMPTY_VALUE);
}
batch.releaseBatch();
}
} catch (IOException e) {
FileIOUtils.deleteFileOrDirectory(localFile);
throw e;
}
return new ContainsFile(localFile, lookupStoreFactory.createReader(localFile));
}
@Override
public void close() throws IOException {
containsFiles.invalidateAll();
}
private static class ContainsFile implements Closeable {
private final File localFile;
private final LookupStoreReader reader;
private boolean isClosed = false;
public ContainsFile(File localFile, LookupStoreReader reader) {
this.localFile = localFile;
this.reader = reader;
}
@Nullable
public byte[] get(byte[] key) throws IOException {
checkArgument(!isClosed);
return reader.lookup(key);
}
@Override
public void close() throws IOException {
reader.close();
isClosed = true;
FileIOUtils.deleteFileOrDirectory(localFile);
}
}
}