How to use

CachingOrcDataSource

in

io.prestosql.orc

@Override
public void readFully(long position, byte[] buffer, int bufferOffset, int length)
    throws IOException
{
  if (position < cachePosition) {
    throw new IllegalArgumentException(String.format("read request (offset %d length %d) is before cache (offset %d length %d)", position, length, cachePosition, cacheLength));
  }
  if (position >= cachePosition + cacheLength) {
    readCacheAt(position);
  }
  if (position + length > cachePosition + cacheLength) {
    throw new IllegalArgumentException(String.format("read request (offset %d length %d) partially overlaps cache (offset %d length %d)", position, length, cachePosition, cacheLength));
  }
  System.arraycopy(cache, toIntExact(position - cachePosition), buffer, bufferOffset, length);
}

@Override
public void readFully(long position, byte[] buffer)
    throws IOException
{
  readFully(position, buffer, 0, buffer.length);
}

CachingOrcDataSource cachingOrcDataSource = new CachingOrcDataSource(
    testingOrcDataSource,
    createTinyStripesRangeFinder(
        maxMergeDistance,
        tinyStripeThreshold));
cachingOrcDataSource.readCacheAt(3);
assertEquals(testingOrcDataSource.getLastReadRanges(), ImmutableList.of(new DiskRange(3, 60)));
cachingOrcDataSource.readCacheAt(63);
assertEquals(testingOrcDataSource.getLastReadRanges(), ImmutableList.of(new DiskRange(63, 8 * 1048576)));
cachingOrcDataSource = new CachingOrcDataSource(
    testingOrcDataSource,
    createTinyStripesRangeFinder(
        maxMergeDistance,
        tinyStripeThreshold));
cachingOrcDataSource.readCacheAt(62); // read at the end of a stripe
assertEquals(testingOrcDataSource.getLastReadRanges(), ImmutableList.of(new DiskRange(3, 60)));
cachingOrcDataSource.readCacheAt(63);
assertEquals(testingOrcDataSource.getLastReadRanges(), ImmutableList.of(new DiskRange(63, 8 * 1048576)));
cachingOrcDataSource = new CachingOrcDataSource(
    testingOrcDataSource,
    createTinyStripesRangeFinder(
        maxMergeDistance,
        tinyStripeThreshold));
cachingOrcDataSource.readCacheAt(3);

@VisibleForTesting
static OrcDataSource wrapWithCacheIfTinyStripes(OrcDataSource dataSource, List<StripeInformation> stripes, DataSize maxMergeDistance, DataSize tinyStripeThreshold)
{
  if (dataSource instanceof CachingOrcDataSource) {
    return dataSource;
  }
  for (StripeInformation stripe : stripes) {
    if (stripe.getTotalLength() > tinyStripeThreshold.toBytes()) {
      return dataSource;
    }
  }
  return new CachingOrcDataSource(dataSource, createTinyStripesRangeFinder(stripes, maxMergeDistance, tinyStripeThreshold));
}

CachingOrcDataSource cachingOrcDataSource = new CachingOrcDataSource(
    testingOrcDataSource,
    createTinyStripesRangeFinder(
        maxMergeDistance,
        tinyStripeThreshold));
cachingOrcDataSource.readCacheAt(3);
assertEquals(testingOrcDataSource.getLastReadRanges(), ImmutableList.of(new DiskRange(3, 60)));
cachingOrcDataSource.readCacheAt(63);
assertEquals(testingOrcDataSource.getLastReadRanges(), ImmutableList.of(new DiskRange(63, 8 * 1048576)));
cachingOrcDataSource = new CachingOrcDataSource(
    testingOrcDataSource,
    createTinyStripesRangeFinder(
        maxMergeDistance,
        tinyStripeThreshold));
cachingOrcDataSource.readCacheAt(62); // read at the end of a stripe
assertEquals(testingOrcDataSource.getLastReadRanges(), ImmutableList.of(new DiskRange(3, 60)));
cachingOrcDataSource.readCacheAt(63);
assertEquals(testingOrcDataSource.getLastReadRanges(), ImmutableList.of(new DiskRange(63, 8 * 1048576)));
cachingOrcDataSource = new CachingOrcDataSource(
    testingOrcDataSource,
    createTinyStripesRangeFinder(
        maxMergeDistance,
        tinyStripeThreshold));
cachingOrcDataSource.readCacheAt(3);

@VisibleForTesting
static OrcDataSource wrapWithCacheIfTinyStripes(OrcDataSource dataSource, List<StripeInformation> stripes, DataSize maxMergeDistance, DataSize tinyStripeThreshold)
{
  if (dataSource instanceof CachingOrcDataSource) {
    return dataSource;
  }
  for (StripeInformation stripe : stripes) {
    if (stripe.getTotalLength() > tinyStripeThreshold.toBytes()) {
      return dataSource;
    }
  }
  return new CachingOrcDataSource(dataSource, createTinyStripesRangeFinder(stripes, maxMergeDistance, tinyStripeThreshold));
}

private static OrcDataSource wrapWithCacheIfTiny(OrcDataSource dataSource, DataSize maxCacheSize)
{
  if (dataSource instanceof CachingOrcDataSource) {
    return dataSource;
  }
  if (dataSource.getSize() > maxCacheSize.toBytes()) {
    return dataSource;
  }
  DiskRange diskRange = new DiskRange(0, toIntExact(dataSource.getSize()));
  return new CachingOrcDataSource(dataSource, desiredOffset -> diskRange);
}

@Override
public void readFully(long position, byte[] buffer, int bufferOffset, int length)
    throws IOException
{
  if (position < cachePosition) {
    throw new IllegalArgumentException(String.format("read request (offset %d length %d) is before cache (offset %d length %d)", position, length, cachePosition, cacheLength));
  }
  if (position >= cachePosition + cacheLength) {
    readCacheAt(position);
  }
  if (position + length > cachePosition + cacheLength) {
    throw new IllegalArgumentException(String.format("read request (offset %d length %d) partially overlaps cache (offset %d length %d)", position, length, cachePosition, cacheLength));
  }
  System.arraycopy(cache, toIntExact(position - cachePosition), buffer, bufferOffset, length);
}

@Override
public void readFully(long position, byte[] buffer)
    throws IOException
{
  readFully(position, buffer, 0, buffer.length);
}

private static OrcDataSource wrapWithCacheIfTiny(OrcDataSource dataSource, DataSize maxCacheSize)
{
  if (dataSource instanceof CachingOrcDataSource) {
    return dataSource;
  }
  if (dataSource.getSize() > maxCacheSize.toBytes()) {
    return dataSource;
  }
  DiskRange diskRange = new DiskRange(0, toIntExact(dataSource.getSize()));
  return new CachingOrcDataSource(dataSource, desiredOffset -> diskRange);
}

@Override
public <K> Map<K, OrcDataSourceInput> readFully(Map<K, DiskRange> diskRanges)
    throws IOException
{
  ImmutableMap.Builder<K, OrcDataSourceInput> builder = ImmutableMap.builder();
  // Assumption here: all disk ranges are in the same region. Therefore, serving them in arbitrary order
  // will not result in eviction of cache that otherwise could have served any of the DiskRanges provided.
  for (Map.Entry<K, DiskRange> entry : diskRanges.entrySet()) {
    DiskRange diskRange = entry.getValue();
    byte[] buffer = new byte[diskRange.getLength()];
    readFully(diskRange.getOffset(), buffer);
    builder.put(entry.getKey(), new OrcDataSourceInput(Slices.wrappedBuffer(buffer).getInput(), buffer.length));
  }
  return builder.build();
}

@Override
public <K> Map<K, OrcDataSourceInput> readFully(Map<K, DiskRange> diskRanges)
    throws IOException
{
  ImmutableMap.Builder<K, OrcDataSourceInput> builder = ImmutableMap.builder();
  // Assumption here: all disk ranges are in the same region. Therefore, serving them in arbitrary order
  // will not result in eviction of cache that otherwise could have served any of the DiskRanges provided.
  for (Map.Entry<K, DiskRange> entry : diskRanges.entrySet()) {
    DiskRange diskRange = entry.getValue();
    byte[] buffer = new byte[diskRange.getLength()];
    readFully(diskRange.getOffset(), buffer);
    builder.put(entry.getKey(), new OrcDataSourceInput(Slices.wrappedBuffer(buffer).getInput(), buffer.length));
  }
  return builder.build();
}