package org.apache.spark.sql.execution.datasources.indexr; import org.apache.hadoop.fs.CreateFlag; import org.apache.hadoop.fs.FSDataOutputStream; import org.apache.hadoop.fs.FileSystem; import org.apache.hadoop.fs.Path; import org.apache.hadoop.fs.permission.FsPermission; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import java.io.IOException; import java.util.EnumSet; import io.indexr.io.ByteBufferWriter; import io.indexr.segment.storage.StorageSegment; import io.indexr.segment.storage.itg.IntegratedSegment; public class SegmentHelper { private static final Logger logger = LoggerFactory.getLogger(SegmentHelper.class); public static long getSegmentBlockSize(FileSystem fileSystem, long fileSize) { long blockSize = fileSystem.getDefaultBlockSize(); while (blockSize < fileSize) { blockSize <<= 1; } return blockSize; } public static void uploadSegment(StorageSegment segment, FileSystem fileSystem, Path path) throws IOException { Path parent = path.getParent(); short replica = fileSystem.getDefaultReplication(parent); if (replica <= 0) { logger.warn("Failed to get replication from {}", parent); replica = fileSystem.getDefaultReplication(); } short _replica = replica; ByteBufferWriter.PredictSizeOpener writeOpener = size -> { long blockSize = getSegmentBlockSize(fileSystem, size); FSDataOutputStream outputStream = fileSystem.create( path, FsPermission.getFileDefault().applyUMask(FsPermission.getUMask(fileSystem.getConf())), EnumSet.of(CreateFlag.CREATE, CreateFlag.OVERWRITE), 4096 * 10, _replica, blockSize, null); ByteBufferWriter writer = ByteBufferWriter.of(outputStream, outputStream::close); writer.setName(path.toString()); return writer; }; IntegratedSegment.Fd.create(segment, writeOpener, null); } public static void notifyUpdate(FileSystem fileSystem, String tableLocation) throws IOException { // Touch the update file to notify segment change. Path updateFilePath = new Path(String.format("%s/__UPDATE__", tableLocation)); try (FSDataOutputStream os = fileSystem.create(updateFilePath, true)) { os.hsync(); } } public static boolean checkSegmentByPath(Path path) { String[] ss = path.toString().split(Path.SEPARATOR); if (ss.length == 0) { return false; } //for (String s : ss) { // if (s.startsWith(".")) { // return false; // } //} String fileName = ss[ss.length - 1]; return !fileName.endsWith("__"); } }