feat(spark): count bytes for files opened with buildColumnarReader(file)

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

@@ -81,9 +81,16 @@ public PartitionReader<ColumnarBatch> createColumnarReader(InputPartition partit
         return VortexReadMetrics.createReader(partition, this, fileOptions, this::buildColumnarReader);
     }
 
+    /**
+     * A reader of one file, for callers that open files themselves rather than through {@link #createColumnarReader}.
+     * Call it on the Spark task thread. It adds the bytes it reads to the task's input metrics, and its
+     * {@code currentMetricsValues} is {@link VortexBytesReadMetric}: one cumulative total of every file opened this way
+     * in the task, so forward the latest reader's values rather than adding them up. Several scans in one task that
+     * open files this way each report their combined total.
+     */
     @Override
     public PartitionReader<ColumnarBatch> buildColumnarReader(PartitionedFile file) {
-        return buildColumnarReader(file, ignored -> {});
+        return VortexReadMetrics.createFileReader(bytesRead -> buildColumnarReader(file, bytesRead));
     }
 
     private PartitionReader<ColumnarBatch> buildColumnarReader(PartitionedFile file, LongConsumer bytesRead) {

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

@@ -76,9 +76,16 @@ public PartitionReader<ColumnarBatch> createColumnarReader(InputPartition partit
         return VortexReadMetrics.createReader(partition, this, fileOptions, this::buildColumnarReader);
     }
 
+    /**
+     * A reader of one file, for callers that open files themselves rather than through {@link #createColumnarReader}.
+     * Call it on the Spark task thread. It adds the bytes it reads to the task's input metrics, and its
+     * {@code currentMetricsValues} is {@link VortexBytesReadMetric}: one cumulative total of every file opened this way
+     * in the task, so forward the latest reader's values rather than adding them up. Several scans in one task that
+     * open files this way each report their combined total.
+     */
     @Override
     public PartitionReader<ColumnarBatch> buildColumnarReader(PartitionedFile file) {
-        return buildColumnarReader(file, ignored -> {});
+        return VortexReadMetrics.createFileReader(bytesRead -> buildColumnarReader(file, bytesRead));
     }
 
     private PartitionReader<ColumnarBatch> buildColumnarReader(PartitionedFile file, LongConsumer bytesRead) {

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

@@ -10,6 +10,7 @@
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.BiFunction;
+import java.util.function.Function;
 import java.util.function.LongConsumer;
 import org.apache.hadoop.fs.FileSystem;
 import org.apache.spark.TaskContext;
@@ -26,6 +27,8 @@
 /** Cumulative successful source bytes, including footers and read-ahead, for one Spark task and scan. */
 final class VortexReadMetrics implements LongConsumer {
     private static final ConcurrentMap<TaskContext, TaskBytes> TASK_BYTES = new ConcurrentHashMap<>();
+    // The scan key of files opened one at a time with buildColumnarReader(file), so they share one task total.
+    private static final Object SINGLE_FILES = new Object();
     // Spark reads this legacy statistics API. A private class key keeps this entry separate from storage statistics.
     @SuppressWarnings("deprecation")
     private static final FileSystem.Statistics NATIVE_STATISTICS =
@@ -113,11 +116,8 @@ static PartitionReader<ColumnarBatch> createReader(
             FileSourceOptions options,
             BiFunction<PartitionedFile, LongConsumer, PartitionReader<ColumnarBatch>> buildReader) {
         TaskContext context = TaskContext.get();
-        TaskBytes taskBytes = taskBytes(context);
-        // Several input partitions can run in one coalesced task. Spark sets (rather than adds) custom metrics,
-        // so each reader of this scan must expose the same cumulative task total.
-        VortexReadMetrics metrics = new VortexReadMetrics(
-                taskBytes.scanBytes.computeIfAbsent(factory, ignored -> new AtomicLong()), taskBytes);
+        // Several input partitions can run in one coalesced task, so each reader of this scan shares one total.
+        VortexReadMetrics metrics = forScan(context, factory);
         PartitionReader<ColumnarBatch> files = new FilePartitionReaderFactory() {
             @Override
             public FileSourceOptions options() {
@@ -136,6 +136,28 @@ public PartitionReader<ColumnarBatch> buildColumnarReader(PartitionedFile file)
         }.createColumnarReader(partition);
 
         // FilePartitionReader does not forward currentMetricsValues from individual files.
+        return reporting(context, files, metrics);
+    }
+
+    /**
+     * A reader of one file, opened outside a scan's {@code createColumnarReader}. Every such file in a task shares one
+     * cumulative total, so a caller that forwards {@code currentMetricsValues} reports all of them.
+     */
+    static PartitionReader<ColumnarBatch> createFileReader(
+            Function<LongConsumer, PartitionReader<ColumnarBatch>> buildReader) {
+        TaskContext context = TaskContext.get();
+        VortexReadMetrics metrics = forScan(context, SINGLE_FILES);
+        return reporting(context, buildReader.apply(metrics), metrics);
+    }
+
+    // Spark sets (rather than adds) custom metrics, so every reader of one scan in a task must share its total.
+    private static VortexReadMetrics forScan(TaskContext context, Object scan) {
+        TaskBytes taskBytes = taskBytes(context);
+        return new VortexReadMetrics(taskBytes.scanBytes.computeIfAbsent(scan, ignored -> new AtomicLong()), taskBytes);
+    }
+
+    private static PartitionReader<ColumnarBatch> reporting(
+            TaskContext context, PartitionReader<ColumnarBatch> files, VortexReadMetrics metrics) {
         return new PartitionReader<>() {
             @Override
             public boolean next() throws IOException {

java/vortex-spark/common/src/test/java/dev/vortex/spark/OneFileAtATimeSource.java

@@ -0,0 +1,200 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+package dev.vortex.spark;
+
+import dev.vortex.spark.io.VortexIo;
+import dev.vortex.spark.read.VortexBytesReadMetric;
+import dev.vortex.spark.read.VortexPartitionReaderFactory;
+import java.io.IOException;
+import java.io.UncheckedIOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.EnumSet;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.stream.Stream;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.spark.paths.SparkPath;
+import org.apache.spark.sql.catalyst.FileSourceOptions;
+import org.apache.spark.sql.catalyst.InternalRow;
+import org.apache.spark.sql.connector.catalog.SupportsRead;
+import org.apache.spark.sql.connector.catalog.Table;
+import org.apache.spark.sql.connector.catalog.TableCapability;
+import org.apache.spark.sql.connector.catalog.TableProvider;
+import org.apache.spark.sql.connector.expressions.Transform;
+import org.apache.spark.sql.connector.metric.CustomMetric;
+import org.apache.spark.sql.connector.metric.CustomTaskMetric;
+import org.apache.spark.sql.connector.read.Batch;
+import org.apache.spark.sql.connector.read.InputPartition;
+import org.apache.spark.sql.connector.read.PartitionReader;
+import org.apache.spark.sql.connector.read.PartitionReaderFactory;
+import org.apache.spark.sql.connector.read.Scan;
+import org.apache.spark.sql.connector.read.ScanBuilder;
+import org.apache.spark.sql.execution.datasources.PartitionedFile;
+import org.apache.spark.sql.sources.Filter;
+import org.apache.spark.sql.types.DataTypes;
+import org.apache.spark.sql.types.StructType;
+import org.apache.spark.sql.util.CaseInsensitiveStringMap;
+import org.apache.spark.sql.vectorized.ColumnarBatch;
+
+/**
+ * A V2 source that reads every Vortex file under {@code path} in one task, with a new reader factory per file, the way
+ * a caller that picks its own files and filters does.
+ */
+public final class OneFileAtATimeSource implements TableProvider {
+    static final StructType SCHEMA = new StructType().add("id", DataTypes.LongType);
+
+    @Override
+    public StructType inferSchema(CaseInsensitiveStringMap options) {
+        return SCHEMA;
+    }
+
+    @Override
+    public Table getTable(StructType schema, Transform[] partitioning, Map<String, String> properties) {
+        List<String> files = vortexFiles(properties.get("path"));
+        return new SourceTable(files);
+    }
+
+    private static List<String> vortexFiles(String directory) {
+        try (Stream<Path> listed = Files.list(Path.of(directory))) {
+            return listed.map(Path::toString)
+                    .filter(f -> f.endsWith(".vortex"))
+                    .sorted()
+                    .toList();
+        } catch (IOException e) {
+            throw new UncheckedIOException(e);
+        }
+    }
+
+    private record SourceTable(List<String> files) implements SupportsRead {
+        @Override
+        public String name() {
+            return "one-file-at-a-time";
+        }
+
+        @Override
+        public StructType schema() {
+            return SCHEMA;
+        }
+
+        @Override
+        public Set<TableCapability> capabilities() {
+            return EnumSet.of(TableCapability.BATCH_READ);
+        }
+
+        @Override
+        public ScanBuilder newScanBuilder(CaseInsensitiveStringMap options) {
+            return () -> new SourceScan(files);
+        }
+    }
+
+    private record SourceScan(List<String> files) implements Scan, Batch {
+        @Override
+        public StructType readSchema() {
+            return SCHEMA;
+        }
+
+        @Override
+        public Batch toBatch() {
+            return this;
+        }
+
+        @Override
+        public CustomMetric[] supportedCustomMetrics() {
+            return new CustomMetric[] {new VortexBytesReadMetric()};
+        }
+
+        @Override
+        public InputPartition[] planInputPartitions() {
+            return new InputPartition[] {new FileList(files)};
+        }
+
+        @Override
+        public PartitionReaderFactory createReaderFactory() {
+            return partition -> new FilesReader(((FileList) partition).paths());
+        }
+    }
+
+    private record FileList(List<String> paths) implements InputPartition {}
+
+    private static final class FilesReader implements PartitionReader<InternalRow> {
+        private final List<String> paths;
+        private final VortexIo io;
+        private final VortexOptions options = VortexOptions.of(Map.of());
+        private int next;
+        private PartitionReader<ColumnarBatch> file;
+        private Iterator<InternalRow> rows;
+
+        private FilesReader(List<String> paths) {
+            this.paths = paths;
+            Configuration conf = new Configuration();
+            conf.set("fs.file.impl", VortexNestedSchemaPruningIoTest.CountingLocalFileSystem.class.getName());
+            conf.set("fs.file.impl.disable.cache", "true");
+            this.io = VortexIo.create(options, conf);
+        }
+
+        @Override
+        public boolean next() throws IOException {
+            while (rows == null || !rows.hasNext()) {
+                if (file != null && file.next()) {
+                    rows = file.get().rowIterator();
+                    continue;
+                }
+                if (next == paths.size()) {
+                    return false;
+                }
+                closeFile();
+                file = open(paths.get(next++));
+                rows = null;
+            }
+            return true;
+        }
+
+        private PartitionReader<ColumnarBatch> open(String path) {
+            long length = Path.of(path).toFile().length();
+            PartitionedFile partitioned = new PartitionedFile(
+                    InternalRow.empty(),
+                    SparkPath.fromPathString("file://" + path),
+                    0,
+                    length,
+                    new String[0],
+                    0,
+                    length,
+                    PartitionedFile.apply$default$8());
+            VortexPartitionReaderFactory factory = new VortexPartitionReaderFactory(
+                    new FileSourceOptions(scala.collection.immutable.Map$.MODULE$.empty()),
+                    io,
+                    options,
+                    SCHEMA,
+                    SCHEMA,
+                    new StructType(),
+                    new Filter[0],
+                    false);
+            return factory.buildColumnarReader(partitioned);
+        }
+
+        @Override
+        public InternalRow get() {
+            return rows.next().copy();
+        }
+
+        @Override
+        public CustomTaskMetric[] currentMetricsValues() {
+            return file == null ? new CustomTaskMetric[0] : file.currentMetricsValues();
+        }
+
+        @Override
+        public void close() throws IOException {
+            closeFile();
+        }
+
+        private void closeFile() throws IOException {
+            if (file != null) {
+                file.close();
+            }
+        }
+    }
+}

java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexInputMetricsTest.java

@@ -99,6 +99,16 @@ void earlyTerminationReportsInputBytes() throws Exception {
         verify(read().limit(1), 1);
     }
 
+    @Test
+    void filesOpenedOneAtATimeReportInputBytes(@TempDir Path tempDir) throws Exception {
+        // A fresh table: Vortex caches what it reads of a file, which would hide reads the other tests already made.
+        String path = tempDir.resolve("one-at-a-time").toString();
+        spark.range(0, 16384, 1, 4).write().format("vortex").save(path);
+        Dataset<Row> read =
+                spark.read().format(OneFileAtATimeSource.class.getName()).load(path);
+        verify(read, 16384);
+    }
+
     private Dataset<Row> read() {
         return spark.read()
                 .schema(SCHEMA)