Report native read bytes in Spark V2 scan metrics

java/vortex-jni/src/main/java/dev/vortex/io/PooledReadable.java

@@ -7,18 +7,20 @@
 import java.io.Closeable;
 import java.io.EOFException;
 import java.io.IOException;
+import java.io.InterruptedIOException;
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Locale;
 import java.util.Queue;
 import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
 
 /**
  * A {@link NativeReadable} that serves each of Vortex's concurrent reads from its own stream.
  *
- * <p>Streams are opened on demand and reused until this readable is closed. Subclasses supply only
- * {@link #openStream()}.
+ * <p>Streams are opened on demand and reused until this readable is closed, which waits for in-flight reads. Subclasses
+ * supply {@link #openStream()}, and may observe successful reads through {@link #onBytesRead(long)}.
  */
 public abstract class PooledReadable implements NativeReadable {
     /** Largest scratch array retained for reuse on a pooled stream; bigger reads get a one-shot array. */
@@ -27,6 +29,7 @@ public abstract class PooledReadable implements NativeReadable {
     private final String name;
     private final long length;
     private final Queue<PooledStream> pool = new ConcurrentLinkedQueue<>();
+    private final ReentrantReadWriteLock lifecycle = new ReentrantReadWriteLock();
     private volatile boolean closed = false;
 
     /**
@@ -43,6 +46,9 @@ protected PooledReadable(String name, long length) {
     /** Opens one more handle on this source. Called whenever a read finds no free stream in the pool. */
     protected abstract PositionalStream openStream() throws IOException;
 
+    /** Called on the reading thread after each read that filled its buffer, with the number of bytes read. */
+    protected void onBytesRead(long bytes) {}
+
     @Override
     public final String name() {
         return name;
@@ -55,6 +61,15 @@ public final long length() {
 
     @Override
     public final void readFully(long position, ByteBuffer buffer) throws IOException {
+        lifecycle.readLock().lock();
+        try {
+            readFullyOpen(position, buffer);
+        } finally {
+            lifecycle.readLock().unlock();
+        }
+    }
+
+    private void readFullyOpen(long position, ByteBuffer buffer) throws IOException {
         Preconditions.checkState(!closed, "Cannot read %s: already closed", name);
         int requested = buffer.remaining();
         if (position < 0 || position + requested > length) {
@@ -99,6 +114,7 @@ public final void readFully(long position, ByteBuffer buffer) throws IOException
             throw failure;
         }
 
+        onBytesRead(requested);
         pool.add(pooled);
         if (closed) {
             // Racing with close(): make sure nothing this read returned to the pool leaks. The read
@@ -114,7 +130,28 @@ public final void readFully(long position, ByteBuffer buffer) throws IOException
     @Override
     public void close() throws IOException {
         closed = true;
-        closeAllPooled();
+        try {
+            // Task cancellation must not wait uninterruptibly behind a stalled storage request, but an
+            // interrupted thread with nothing in flight should still close cleanly.
+            if (!lifecycle.writeLock().tryLock()) {
+                lifecycle.writeLock().lockInterruptibly();
+            }
+        } catch (InterruptedException e) {
+            Thread.currentThread().interrupt();
+            InterruptedIOException failure = new InterruptedIOException("Interrupted closing " + name);
+            failure.initCause(e);
+            try {
+                closeAllPooled();
+            } catch (IOException closeFailure) {
+                failure.addSuppressed(closeFailure);
+            }
+            throw failure;
+        }
+        try {
+            closeAllPooled();
+        } finally {
+            lifecycle.writeLock().unlock();
+        }
     }
 
     private void closeAllPooled() throws IOException {

java/vortex-spark/common/src/main/java/dev/vortex/spark/io/HadoopReadable.java

@@ -8,8 +8,10 @@
 import java.io.InterruptedIOException;
 import java.io.UncheckedIOException;
 import java.nio.ByteBuffer;
+import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
+import java.util.function.LongConsumer;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.fs.FSDataInputStream;
 import org.apache.hadoop.fs.FileSystem;
@@ -24,6 +26,7 @@ public final class HadoopReadable extends PooledReadable {
     private final FileSystem fs;
     private final int coalesceDistance;
     private final int coalesceMaxSize;
+    private final LongConsumer bytesRead;
 
     /** Opens a readable over {@code path}, stating the file for its length. */
     public static HadoopReadable open(Configuration conf, String path) {
@@ -35,7 +38,7 @@ public static HadoopReadable open(Configuration conf, String path) {
      * stats every file when it plans the scan, so executors need not stat them again.
      */
     public static HadoopReadable open(Configuration conf, String path, long length) {
-        return open(conf, path, length, DEFAULT_COALESCE_DISTANCE, DEFAULT_COALESCE_MAX_SIZE);
+        return open(conf, path, length, DEFAULT_COALESCE_DISTANCE, DEFAULT_COALESCE_MAX_SIZE, ignored -> {});
     }
 
     /**
@@ -44,23 +47,41 @@ public static HadoopReadable open(Configuration conf, String path, long length)
      */
     public static HadoopReadable open(
             Configuration conf, String path, long length, int coalesceDistance, int coalesceMaxSize) {
+        return open(conf, path, length, coalesceDistance, coalesceMaxSize, ignored -> {});
+    }
+
+    /** As {@link #open(Configuration, String, long, int, int)}, reporting each successful read's byte count. */
+    public static HadoopReadable open(
+            Configuration conf,
+            String path,
+            long length,
+            int coalesceDistance,
+            int coalesceMaxSize,
+            LongConsumer bytesRead) {
         Path file = new Path(path);
         try {
             FileSystem fs = file.getFileSystem(conf);
             long resolved = length >= 0 ? length : fs.getFileStatus(file).getLen();
-            return new HadoopReadable(path, file, fs, resolved, coalesceDistance, coalesceMaxSize);
+            return new HadoopReadable(path, file, fs, resolved, coalesceDistance, coalesceMaxSize, bytesRead);
         } catch (IOException e) {
             throw new UncheckedIOException("Failed to open " + path, e);
         }
     }
 
     private HadoopReadable(
-            String name, Path path, FileSystem fs, long length, int coalesceDistance, int coalesceMaxSize) {
+            String name,
+            Path path,
+            FileSystem fs,
+            long length,
+            int coalesceDistance,
+            int coalesceMaxSize,
+            LongConsumer bytesRead) {
         super(name, length);
         this.path = path;
         this.fs = fs;
         this.coalesceDistance = coalesceDistance;
         this.coalesceMaxSize = coalesceMaxSize;
+        this.bytesRead = Objects.requireNonNull(bytesRead, "bytesRead");
     }
 
     @Override
@@ -95,6 +116,11 @@ protected PositionalStream openStream() throws IOException {
         }
     }
 
+    @Override
+    protected void onBytesRead(long bytes) {
+        bytesRead.accept(bytes);
+    }
+
     private static final class HadoopStream implements PositionalStream {
         private final FSDataInputStream stream;
         // Hadoop mandates this capability probe: streams may implement ByteBufferPositionedReadable

java/vortex-spark/common/src/main/java/dev/vortex/spark/io/VortexIo.java

@@ -8,6 +8,7 @@
 import java.io.Serializable;
 import java.util.Locale;
 import java.util.Objects;
+import java.util.function.LongConsumer;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.spark.SerializableWritable;
 import org.apache.spark.SparkContext;
@@ -163,8 +164,18 @@ public int splitConcurrency() {
      * @throws IllegalArgumentException if the path does not carry {@link VortexFile#EXTENSION}
      */
     public NativeReadable openReadable(VortexFile file) {
+        return openReadable(file, ignored -> {});
+    }
+
+    /** Opens a source that reports each successful read's byte count, on the thread performing the read. */
+    public NativeReadable openReadable(VortexFile file, LongConsumer bytesRead) {
         VortexFile.requireVortexExtension(file.path());
         return HadoopReadable.open(
-                hadoopConf(), file.path(), file.length(), settings.coalesceDistance(), settings.coalesceMaxSize());
+                hadoopConf(),
+                file.path(),
+                file.length(),
+                settings.coalesceDistance(),
+                settings.coalesceMaxSize(),
+                bytesRead);
     }
 }

java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexAggregateReaderFactory.java

@@ -9,6 +9,7 @@
 import dev.vortex.spark.io.VortexIo;
 import java.io.Serializable;
 import java.util.OptionalLong;
+import java.util.function.LongConsumer;
 import org.apache.spark.sql.catalyst.FileSourceOptions;
 import org.apache.spark.sql.catalyst.InternalRow;
 import org.apache.spark.sql.connector.expressions.Expression;
@@ -75,8 +76,17 @@ public PartitionReader<InternalRow> buildReader(PartitionedFile file) {
         throw new UnsupportedOperationException("row-based aggregate reads are not supported");
     }
 
+    @Override
+    public PartitionReader<ColumnarBatch> createColumnarReader(InputPartition partition) {
+        return VortexReadMetrics.createReader(partition, this, fileOptions, this::buildColumnarReader);
+    }
+
     @Override
     public PartitionReader<ColumnarBatch> buildColumnarReader(PartitionedFile file) {
+        return buildColumnarReader(file, ignored -> {});
+    }
+
+    private PartitionReader<ColumnarBatch> buildColumnarReader(PartitionedFile file, LongConsumer bytesRead) {
         if (!VortexFile.hasVortexExtension(file.toPath().getName())) {
             return new EmptyPartitionReader<>();
         }
@@ -93,7 +103,7 @@ public boolean next() {
                 emitted = true;
                 // An estimate is no answer to COUNT(*), so the footer must state the count exactly.
                 OptionalLong rowCount = VortexFooterReader.exactRowCount(
-                        new VortexFile(file.toPath().toString(), file.fileSize()), io, formatOptions);
+                        new VortexFile(file.toPath().toString(), file.fileSize()), io, formatOptions, bytesRead);
                 if (rowCount.isEmpty()) {
                     throw new IllegalStateException(String.format(
                             "Vortex footer states no exact row count for %s, so COUNT(*) cannot be answered from it. "

java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexBytesReadMetric.java

@@ -0,0 +1,21 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+package dev.vortex.spark.read;
+
+import org.apache.spark.sql.connector.metric.CustomSumMetric;
+
+/** Source bytes successfully read by Vortex, including footers and read-ahead. */
+public final class VortexBytesReadMetric extends CustomSumMetric {
+    static final String NAME = "vortexBytesRead";
+
+    @Override
+    public String name() {
+        return NAME;
+    }
+
+    @Override
+    public String description() {
+        return "Vortex source bytes read";
+    }
+}

java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexFooterReader.java

@@ -27,6 +27,7 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
+import java.util.function.LongConsumer;
 import org.apache.hadoop.fs.FileStatus;
 import org.apache.spark.sql.types.ArrayType;
 import org.apache.spark.sql.types.DataType;
@@ -132,10 +133,15 @@ public static OptionalLong estimatedRowCount(VortexFile file, VortexIo io, Vorte
      * be returned to the user as fact.
      */
     public static OptionalLong exactRowCount(VortexFile file, VortexIo io, VortexOptions options) {
+        return exactRowCount(file, io, options, ignored -> {});
+    }
+
+    static OptionalLong exactRowCount(VortexFile file, VortexIo io, VortexOptions options, LongConsumer bytesRead) {
         return withDataSource(
                 file,
                 io,
                 options,
+                bytesRead,
                 source -> source.rowCount() instanceof RowCount.Exact exact
                         ? OptionalLong.of(exact.value())
                         : OptionalLong.empty());
@@ -328,8 +334,17 @@ private static <T> List<T> mapFooters(
 
     private static <T> T withDataSource(
             VortexFile file, VortexIo io, VortexOptions options, DataSourceFunction<T> function) {
+        return withDataSource(file, io, options, ignored -> {}, function);
+    }
+
+    private static <T> T withDataSource(
+            VortexFile file,
+            VortexIo io,
+            VortexOptions options,
+            LongConsumer bytesRead,
+            DataSourceFunction<T> function) {
         Session session = VortexSparkSession.get(options);
-        try (NativeReadable readable = io.openReadable(file);
+        try (NativeReadable readable = io.openReadable(file, bytesRead);
                 DataSource source = DataSource.open(session, List.of(readable), io.readConcurrency())) {
             return function.apply(source);
         } catch (IOException e) {

java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexPartitionReader.java

@@ -30,6 +30,7 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.TreeMap;
+import java.util.function.LongConsumer;
 import org.apache.spark.sql.connector.read.PartitionReader;
 import org.apache.spark.sql.execution.datasources.PartitionedFile;
 import org.apache.spark.sql.execution.vectorized.ConstantColumnVector;
@@ -78,13 +79,35 @@ public VortexPartitionReader(
             VortexOptions formatOptions,
             Filter[] pushedFilters,
             boolean caseSensitive) {
+        this(
+                file,
+                dataSchema,
+                readDataSchema,
+                readPartitionSchema,
+                io,
+                formatOptions,
+                pushedFilters,
+                caseSensitive,
+                ignored -> {});
+    }
+
+    VortexPartitionReader(
+            PartitionedFile file,
+            StructType dataSchema,
+            StructType readDataSchema,
+            StructType readPartitionSchema,
+            VortexIo io,
+            VortexOptions formatOptions,
+            Filter[] pushedFilters,
+            boolean caseSensitive,
+            LongConsumer bytesRead) {
         this.file = file;
         this.readDataSchema = readDataSchema;
         this.readPartitionSchema = readPartitionSchema;
         this.allocator = ArrowAllocation.rootAllocator();
         try {
             session = VortexSparkSession.get(formatOptions);
-            readable = io.openReadable(new VortexFile(file.toPath().toString(), file.fileSize()));
+            readable = io.openReadable(new VortexFile(file.toPath().toString(), file.fileSize()), bytesRead);
             dataSource = DataSource.open(session, List.of(readable), io.readConcurrency());
 
             List<Field> arrowFields = dataSource.arrowSchema(allocator).getFields();

java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexPartitionReaderFactory.java

@@ -8,6 +8,7 @@
 import dev.vortex.spark.io.VortexFile;
 import dev.vortex.spark.io.VortexIo;
 import java.io.Serializable;
+import java.util.function.LongConsumer;
 import org.apache.spark.sql.catalyst.FileSourceOptions;
 import org.apache.spark.sql.catalyst.InternalRow;
 import org.apache.spark.sql.connector.read.InputPartition;
@@ -70,14 +71,31 @@ public PartitionReader<InternalRow> buildReader(PartitionedFile file) {
         throw new UnsupportedOperationException("row-based V2 reads are not supported");
     }
 
+    @Override
+    public PartitionReader<ColumnarBatch> createColumnarReader(InputPartition partition) {
+        return VortexReadMetrics.createReader(partition, this, fileOptions, this::buildColumnarReader);
+    }
+
     @Override
     public PartitionReader<ColumnarBatch> buildColumnarReader(PartitionedFile file) {
+        return buildColumnarReader(file, ignored -> {});
+    }
+
+    private PartitionReader<ColumnarBatch> buildColumnarReader(PartitionedFile file, LongConsumer bytesRead) {
         if (!VortexFile.hasVortexExtension(file.toPath().getName())) {
             return new EmptyPartitionReader<>();
         }
 
         return new VortexPartitionReader(
-                file, dataSchema, readDataSchema, readPartitionSchema, io, formatOptions, pushedFilters, caseSensitive);
+                file,
+                dataSchema,
+                readDataSchema,
+                readPartitionSchema,
+                io,
+                formatOptions,
+                pushedFilters,
+                caseSensitive,
+                bytesRead);
     }
 
     @Override