From d34fe7d75464ad3d1b79f7fcd973f962050adc5e Mon Sep 17 00:00:00 2001 From: Christopher Wirt Date: Tue, 10 Nov 2026 15:25:50 +0000 Subject: [PATCH 1/3] feat(spark): let readers that open their own files report bytes read #108 reports native read bytes to Spark, but only through the connector's own scan: VortexReadMetrics is package-private, and the factory's reader that takes a byte counter is private. A data source that builds VortexPartitionReaderFactory itself and reads files one by one, as the event store's history read does, reports 0 input bytes. - VortexReadMetrics is public, with forTask(scan) for a task's cumulative counter, publish() to move worker-thread bytes into the task's input metrics, and currentMetricsValues(). - VortexPartitionReaderFactory.buildColumnarReader(file, metrics) is public. - The connector's own scan gets its counter from forTask, so both paths share one implementation. Co-Authored-By: Claude Opus 5.5 --- .../read/VortexPartitionReaderFactory.java | 11 ++- .../vortex/spark/read/VortexReadMetrics.java | 54 ++++++++---- .../vortex/spark/VortexInputMetricsTest.java | 88 +++++++++++++++++++ 3 files changed, 137 insertions(+), 16 deletions(-) diff --git a/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexPartitionReaderFactory.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexPartitionReaderFactory.java index 5091c878869..855e47ce85f 100644 --- a/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexPartitionReaderFactory.java +++ b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexPartitionReaderFactory.java @@ -73,7 +73,8 @@ public PartitionReader buildReader(PartitionedFile file) { @Override public PartitionReader createColumnarReader(InputPartition partition) { - return VortexReadMetrics.createReader(partition, this, fileOptions, this::buildColumnarReader); + return VortexReadMetrics.createReader( + partition, this, fileOptions, (file, metrics) -> buildColumnarReader(file, (LongConsumer) metrics)); } @Override @@ -81,6 +82,14 @@ public PartitionReader buildColumnarReader(PartitionedFile file) return buildColumnarReader(file, ignored -> {}); } + /** + * A reader of {@code file} that counts its source bytes in {@code metrics}, for callers that open files one by one + * rather than through {@link #createColumnarReader}. + */ + public PartitionReader buildColumnarReader(PartitionedFile file, VortexReadMetrics metrics) { + return buildColumnarReader(file, (LongConsumer) metrics); + } + private PartitionReader buildColumnarReader(PartitionedFile file, LongConsumer bytesRead) { if (!VortexFile.hasVortexExtension(file.toPath().getName())) { return new EmptyPartitionReader<>(); diff --git a/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexReadMetrics.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexReadMetrics.java index 7ae5f08a2cf..a1365345c86 100644 --- a/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexReadMetrics.java +++ b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexReadMetrics.java @@ -23,8 +23,15 @@ import org.apache.spark.sql.vectorized.ColumnarBatch; import org.apache.spark.util.TaskCompletionListener; -/** Cumulative successful source bytes, including footers and read-ahead, for one Spark task and scan. */ -final class VortexReadMetrics implements LongConsumer { +/** + * Cumulative successful source bytes, including footers and read-ahead, for one Spark task and scan. + * + *

The connector's own scans count their bytes with it. A reader that opens files itself, with + * {@link VortexPartitionReaderFactory#buildColumnarReader(PartitionedFile, VortexReadMetrics)}, gets one from + * {@link #forTask}, reports {@link #currentMetricsValues} as its own, and calls {@link #publish} on the task thread + * after each file, so Spark's input bytes include the reads Vortex makes on its worker threads. + */ +public final class VortexReadMetrics implements LongConsumer { private static final ConcurrentMap TASK_BYTES = new ConcurrentHashMap<>(); // Spark reads this legacy statistics API. A private class key keeps this entry separate from storage statistics. @SuppressWarnings("deprecation") @@ -35,15 +42,29 @@ private abstract static class NativeReadStatistics extends FileSystem {} private final AtomicLong bytesRead; private final TaskBytes taskBytes; + private final TaskContext context; private final Thread taskThread = Thread.currentThread(); VortexReadMetrics() { - this(new AtomicLong(), new TaskBytes()); + this(new AtomicLong(), new TaskBytes(), null); } - private VortexReadMetrics(AtomicLong bytesRead, TaskBytes taskBytes) { + private VortexReadMetrics(AtomicLong bytesRead, TaskBytes taskBytes, TaskContext context) { this.bytesRead = bytesRead; this.taskBytes = taskBytes; + this.context = context; + } + + /** + * The metrics of {@code scan}'s reads in the current Spark task, called on the task thread. Every call with the + * same {@code scan} in one task shares one cumulative total, as Spark sets custom metrics rather than adding them. + * Outside a task, the bytes are counted but published nowhere. + */ + public static VortexReadMetrics forTask(Object scan) { + TaskContext context = TaskContext.get(); + TaskBytes taskBytes = taskBytes(context); + return new VortexReadMetrics( + taskBytes.scanBytes.computeIfAbsent(scan, ignored -> new AtomicLong()), taskBytes, context); } @Override @@ -68,7 +89,11 @@ private static TaskBytes taskBytes(TaskContext context) { return created; } - private void publishWorkerBytes(TaskContext context) { + /** + * Adds the bytes read on Vortex's worker threads since the last call to the task's input metrics. Call it on the + * task thread, after a file's reader is closed so its read-ahead is counted. + */ + public void publish() { if (context == null) { return; } @@ -90,7 +115,8 @@ private static final class TaskBytes { private long publishedWorkerBytes; } - CustomTaskMetric[] currentMetricsValues() { + /** A snapshot of the cumulative total, as {@link VortexBytesReadMetric}. */ + public CustomTaskMetric[] currentMetricsValues() { long snapshot = bytesRead.get(); return new CustomTaskMetric[] { new CustomTaskMetric() { @@ -111,13 +137,11 @@ static PartitionReader createReader( InputPartition partition, Object factory, FileSourceOptions options, - BiFunction> buildReader) { + BiFunction> 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 exposes the same + // cumulative task total. + VortexReadMetrics metrics = forTask(factory); PartitionReader files = new FilePartitionReaderFactory() { @Override public FileSourceOptions options() { @@ -142,7 +166,7 @@ public boolean next() throws IOException { try { return files.next(); } finally { - metrics.publishWorkerBytes(context); + metrics.publish(); } } @@ -161,7 +185,7 @@ public CustomTaskMetric[] currentMetricsValues() { throw new UncheckedIOException(e); } } - metrics.publishWorkerBytes(context); + metrics.publish(); return metrics.currentMetricsValues(); } @@ -170,7 +194,7 @@ public void close() throws IOException { try { files.close(); } finally { - metrics.publishWorkerBytes(context); + metrics.publish(); } } }; diff --git a/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexInputMetricsTest.java b/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexInputMetricsTest.java index 4a65a3b6f24..d444de2724c 100644 --- a/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexInputMetricsTest.java +++ b/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexInputMetricsTest.java @@ -6,16 +6,34 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertTrue; +import dev.vortex.spark.io.VortexIo; +import dev.vortex.spark.read.VortexPartitionReaderFactory; +import dev.vortex.spark.read.VortexReadMetrics; +import java.io.IOException; +import java.io.UncheckedIOException; +import java.nio.file.Files; import java.nio.file.Path; +import java.util.List; +import java.util.Map; import java.util.concurrent.atomic.AtomicLong; +import java.util.stream.Stream; +import org.apache.hadoop.conf.Configuration; +import org.apache.spark.api.java.JavaSparkContext; +import org.apache.spark.paths.SparkPath; import org.apache.spark.scheduler.SparkListener; import org.apache.spark.scheduler.SparkListenerTaskEnd; import org.apache.spark.sql.Dataset; import org.apache.spark.sql.Row; import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.catalyst.FileSourceOptions; +import org.apache.spark.sql.catalyst.InternalRow; +import org.apache.spark.sql.connector.read.PartitionReader; import org.apache.spark.sql.execution.SparkPlan; +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.vectorized.ColumnarBatch; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; @@ -99,6 +117,76 @@ void earlyTerminationReportsInputBytes() throws Exception { verify(read().limit(1), 1); } + @Test + void filesReadOneByOneThroughTheFactoryReportTheirBytes() throws Exception { + // Its own files: Vortex caches what it reads of a file, which would hide the footer reads other tests count. + String own = table + "-one-by-one"; + spark.range(0, 16384, 1, 4).write().format("vortex").save(own); + List files; + try (Stream listed = Files.list(Path.of(own))) { + files = listed.map(Path::toString) + .filter(f -> f.endsWith(".vortex")) + .sorted() + .toList(); + } + assertTrue(files.size() > 1, "the table must have several files"); + spark.sparkContext().listenerBus().waitUntilEmpty(10000); + listener.bytes.set(0); + VortexNestedSchemaPruningIoTest.BYTES_READ.set(0); + + List results = JavaSparkContext.fromSparkContext(spark.sparkContext()) + .parallelize(List.of(0), 1) + .map(ignored -> readOneByOne(files)) + .collect(); + spark.sparkContext().listenerBus().waitUntilEmpty(10000); + + long source = VortexNestedSchemaPruningIoTest.BYTES_READ.get(); + assertEquals(16384, results.get(0)[0]); + assertTrue(source > 0, "the reads must read source bytes"); + assertEquals(source, results.get(0)[1], "the custom metric must hold every file's bytes"); + assertEquals(source, listener.bytes.get(), "task metrics must include each source byte once"); + } + + /** Reads each file with its own reader in this task, as a caller choosing its own files would: {rows, bytes}. */ + private static long[] readOneByOne(List files) { + VortexOptions options = VortexOptions.of(Map.of()); + Configuration conf = new Configuration(); + conf.set("fs.file.impl", VortexNestedSchemaPruningIoTest.CountingLocalFileSystem.class.getName()); + conf.set("fs.file.impl.disable.cache", "true"); + VortexPartitionReaderFactory factory = new VortexPartitionReaderFactory( + new FileSourceOptions(scala.collection.immutable.Map$.MODULE$.empty()), + VortexIo.create(options, conf), + options, + SCHEMA, + SCHEMA, + new StructType(), + new Filter[0], + false); + VortexReadMetrics metrics = VortexReadMetrics.forTask(factory); + long rows = 0; + for (String path : files) { + long length = Path.of(path).toFile().length(); + PartitionedFile file = new PartitionedFile( + InternalRow.empty(), + SparkPath.fromPathString("file://" + path), + 0, + length, + new String[0], + 0, + length, + PartitionedFile.apply$default$8()); + try (PartitionReader reader = factory.buildColumnarReader(file, metrics)) { + while (reader.next()) { + rows += reader.get().numRows(); + } + } catch (IOException e) { + throw new UncheckedIOException(e); + } + metrics.publish(); + } + return new long[] {rows, metrics.currentMetricsValues()[0].value()}; + } + private Dataset read() { return spark.read() .schema(SCHEMA) From ec92d56ef89456d1d1da17fd4bd63501bf47eb14 Mon Sep 17 00:00:00 2001 From: Christopher Wirt Date: Wed, 11 Nov 2026 10:45:17 +0000 Subject: [PATCH 2/3] feat(spark): count bytes for files opened with buildColumnarReader(file) Replaces the public VortexReadMetrics API with self-reporting readers. The public buildColumnarReader(file) now counts its bytes, publishes them on next() and close(), and exposes a cumulative vortexBytesRead shared by every such file in the task. Callers that open files one by one need no changes to report input bytes. Co-Authored-By: Claude Opus 5.5 --- .../read/VortexAggregateReaderFactory.java | 2 +- .../read/VortexPartitionReaderFactory.java | 13 +- .../vortex/spark/read/VortexReadMetrics.java | 77 ++++--- .../vortex/spark/OneFileAtATimeSource.java | 200 ++++++++++++++++++ .../vortex/spark/VortexInputMetricsTest.java | 92 +------- 5 files changed, 247 insertions(+), 137 deletions(-) create mode 100644 java/vortex-spark/common/src/test/java/dev/vortex/spark/OneFileAtATimeSource.java diff --git a/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexAggregateReaderFactory.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexAggregateReaderFactory.java index ddf0b8bafa5..2bfa2bd09de 100644 --- a/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexAggregateReaderFactory.java +++ b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexAggregateReaderFactory.java @@ -83,7 +83,7 @@ public PartitionReader createColumnarReader(InputPartition partit @Override public PartitionReader buildColumnarReader(PartitionedFile file) { - return buildColumnarReader(file, ignored -> {}); + return VortexReadMetrics.createFileReader(bytesRead -> buildColumnarReader(file, bytesRead)); } private PartitionReader buildColumnarReader(PartitionedFile file, LongConsumer bytesRead) { diff --git a/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexPartitionReaderFactory.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexPartitionReaderFactory.java index 855e47ce85f..638fe38896e 100644 --- a/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexPartitionReaderFactory.java +++ b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexPartitionReaderFactory.java @@ -73,21 +73,12 @@ public PartitionReader buildReader(PartitionedFile file) { @Override public PartitionReader createColumnarReader(InputPartition partition) { - return VortexReadMetrics.createReader( - partition, this, fileOptions, (file, metrics) -> buildColumnarReader(file, (LongConsumer) metrics)); + return VortexReadMetrics.createReader(partition, this, fileOptions, this::buildColumnarReader); } @Override public PartitionReader buildColumnarReader(PartitionedFile file) { - return buildColumnarReader(file, ignored -> {}); - } - - /** - * A reader of {@code file} that counts its source bytes in {@code metrics}, for callers that open files one by one - * rather than through {@link #createColumnarReader}. - */ - public PartitionReader buildColumnarReader(PartitionedFile file, VortexReadMetrics metrics) { - return buildColumnarReader(file, (LongConsumer) metrics); + return VortexReadMetrics.createFileReader(bytesRead -> buildColumnarReader(file, bytesRead)); } private PartitionReader buildColumnarReader(PartitionedFile file, LongConsumer bytesRead) { diff --git a/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexReadMetrics.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexReadMetrics.java index a1365345c86..70f5c587f56 100644 --- a/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexReadMetrics.java +++ b/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; @@ -23,16 +24,11 @@ import org.apache.spark.sql.vectorized.ColumnarBatch; import org.apache.spark.util.TaskCompletionListener; -/** - * Cumulative successful source bytes, including footers and read-ahead, for one Spark task and scan. - * - *

The connector's own scans count their bytes with it. A reader that opens files itself, with - * {@link VortexPartitionReaderFactory#buildColumnarReader(PartitionedFile, VortexReadMetrics)}, gets one from - * {@link #forTask}, reports {@link #currentMetricsValues} as its own, and calls {@link #publish} on the task thread - * after each file, so Spark's input bytes include the reads Vortex makes on its worker threads. - */ -public final class VortexReadMetrics implements LongConsumer { +/** Cumulative successful source bytes, including footers and read-ahead, for one Spark task and scan. */ +final class VortexReadMetrics implements LongConsumer { private static final ConcurrentMap 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 = @@ -42,29 +38,15 @@ private abstract static class NativeReadStatistics extends FileSystem {} private final AtomicLong bytesRead; private final TaskBytes taskBytes; - private final TaskContext context; private final Thread taskThread = Thread.currentThread(); VortexReadMetrics() { - this(new AtomicLong(), new TaskBytes(), null); + this(new AtomicLong(), new TaskBytes()); } - private VortexReadMetrics(AtomicLong bytesRead, TaskBytes taskBytes, TaskContext context) { + private VortexReadMetrics(AtomicLong bytesRead, TaskBytes taskBytes) { this.bytesRead = bytesRead; this.taskBytes = taskBytes; - this.context = context; - } - - /** - * The metrics of {@code scan}'s reads in the current Spark task, called on the task thread. Every call with the - * same {@code scan} in one task shares one cumulative total, as Spark sets custom metrics rather than adding them. - * Outside a task, the bytes are counted but published nowhere. - */ - public static VortexReadMetrics forTask(Object scan) { - TaskContext context = TaskContext.get(); - TaskBytes taskBytes = taskBytes(context); - return new VortexReadMetrics( - taskBytes.scanBytes.computeIfAbsent(scan, ignored -> new AtomicLong()), taskBytes, context); } @Override @@ -89,11 +71,7 @@ private static TaskBytes taskBytes(TaskContext context) { return created; } - /** - * Adds the bytes read on Vortex's worker threads since the last call to the task's input metrics. Call it on the - * task thread, after a file's reader is closed so its read-ahead is counted. - */ - public void publish() { + private void publishWorkerBytes(TaskContext context) { if (context == null) { return; } @@ -115,8 +93,7 @@ private static final class TaskBytes { private long publishedWorkerBytes; } - /** A snapshot of the cumulative total, as {@link VortexBytesReadMetric}. */ - public CustomTaskMetric[] currentMetricsValues() { + CustomTaskMetric[] currentMetricsValues() { long snapshot = bytesRead.get(); return new CustomTaskMetric[] { new CustomTaskMetric() { @@ -137,11 +114,9 @@ static PartitionReader createReader( InputPartition partition, Object factory, FileSourceOptions options, - BiFunction> buildReader) { - TaskContext context = TaskContext.get(); - // Several input partitions can run in one coalesced task, so each reader of this scan exposes the same - // cumulative task total. - VortexReadMetrics metrics = forTask(factory); + BiFunction> buildReader) { + // Several input partitions can run in one coalesced task, so each reader of this scan shares one total. + VortexReadMetrics metrics = forScan(factory); PartitionReader files = new FilePartitionReaderFactory() { @Override public FileSourceOptions options() { @@ -160,13 +135,35 @@ public PartitionReader buildColumnarReader(PartitionedFile file) }.createColumnarReader(partition); // FilePartitionReader does not forward currentMetricsValues from individual files. + return reporting(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 createFileReader( + Function> buildReader) { + VortexReadMetrics metrics = forScan(SINGLE_FILES); + return reporting(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(Object scan) { + TaskBytes taskBytes = taskBytes(TaskContext.get()); + return new VortexReadMetrics(taskBytes.scanBytes.computeIfAbsent(scan, ignored -> new AtomicLong()), taskBytes); + } + + private static PartitionReader reporting( + PartitionReader files, VortexReadMetrics metrics) { + TaskContext context = TaskContext.get(); return new PartitionReader<>() { @Override public boolean next() throws IOException { try { return files.next(); } finally { - metrics.publish(); + metrics.publishWorkerBytes(context); } } @@ -185,7 +182,7 @@ public CustomTaskMetric[] currentMetricsValues() { throw new UncheckedIOException(e); } } - metrics.publish(); + metrics.publishWorkerBytes(context); return metrics.currentMetricsValues(); } @@ -194,7 +191,7 @@ public void close() throws IOException { try { files.close(); } finally { - metrics.publish(); + metrics.publishWorkerBytes(context); } } }; diff --git a/java/vortex-spark/common/src/test/java/dev/vortex/spark/OneFileAtATimeSource.java b/java/vortex-spark/common/src/test/java/dev/vortex/spark/OneFileAtATimeSource.java new file mode 100644 index 00000000000..29e5ef6994f --- /dev/null +++ b/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 properties) { + List files = vortexFiles(properties.get("path")); + return new SourceTable(files); + } + + private static List vortexFiles(String directory) { + try (Stream 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 files) implements SupportsRead { + @Override + public String name() { + return "one-file-at-a-time"; + } + + @Override + public StructType schema() { + return SCHEMA; + } + + @Override + public Set capabilities() { + return EnumSet.of(TableCapability.BATCH_READ); + } + + @Override + public ScanBuilder newScanBuilder(CaseInsensitiveStringMap options) { + return () -> new SourceScan(files); + } + } + + private record SourceScan(List 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 paths) implements InputPartition {} + + private static final class FilesReader implements PartitionReader { + private final List paths; + private final VortexIo io; + private final VortexOptions options = VortexOptions.of(Map.of()); + private int next; + private PartitionReader file; + private Iterator rows; + + private FilesReader(List 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 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(); + } + } + } +} diff --git a/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexInputMetricsTest.java b/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexInputMetricsTest.java index d444de2724c..e34845f3aa5 100644 --- a/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexInputMetricsTest.java +++ b/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexInputMetricsTest.java @@ -6,34 +6,16 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertTrue; -import dev.vortex.spark.io.VortexIo; -import dev.vortex.spark.read.VortexPartitionReaderFactory; -import dev.vortex.spark.read.VortexReadMetrics; -import java.io.IOException; -import java.io.UncheckedIOException; -import java.nio.file.Files; import java.nio.file.Path; -import java.util.List; -import java.util.Map; import java.util.concurrent.atomic.AtomicLong; -import java.util.stream.Stream; -import org.apache.hadoop.conf.Configuration; -import org.apache.spark.api.java.JavaSparkContext; -import org.apache.spark.paths.SparkPath; import org.apache.spark.scheduler.SparkListener; import org.apache.spark.scheduler.SparkListenerTaskEnd; import org.apache.spark.sql.Dataset; import org.apache.spark.sql.Row; import org.apache.spark.sql.SparkSession; -import org.apache.spark.sql.catalyst.FileSourceOptions; -import org.apache.spark.sql.catalyst.InternalRow; -import org.apache.spark.sql.connector.read.PartitionReader; import org.apache.spark.sql.execution.SparkPlan; -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.vectorized.ColumnarBatch; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; @@ -118,73 +100,13 @@ void earlyTerminationReportsInputBytes() throws Exception { } @Test - void filesReadOneByOneThroughTheFactoryReportTheirBytes() throws Exception { - // Its own files: Vortex caches what it reads of a file, which would hide the footer reads other tests count. - String own = table + "-one-by-one"; - spark.range(0, 16384, 1, 4).write().format("vortex").save(own); - List files; - try (Stream listed = Files.list(Path.of(own))) { - files = listed.map(Path::toString) - .filter(f -> f.endsWith(".vortex")) - .sorted() - .toList(); - } - assertTrue(files.size() > 1, "the table must have several files"); - spark.sparkContext().listenerBus().waitUntilEmpty(10000); - listener.bytes.set(0); - VortexNestedSchemaPruningIoTest.BYTES_READ.set(0); - - List results = JavaSparkContext.fromSparkContext(spark.sparkContext()) - .parallelize(List.of(0), 1) - .map(ignored -> readOneByOne(files)) - .collect(); - spark.sparkContext().listenerBus().waitUntilEmpty(10000); - - long source = VortexNestedSchemaPruningIoTest.BYTES_READ.get(); - assertEquals(16384, results.get(0)[0]); - assertTrue(source > 0, "the reads must read source bytes"); - assertEquals(source, results.get(0)[1], "the custom metric must hold every file's bytes"); - assertEquals(source, listener.bytes.get(), "task metrics must include each source byte once"); - } - - /** Reads each file with its own reader in this task, as a caller choosing its own files would: {rows, bytes}. */ - private static long[] readOneByOne(List files) { - VortexOptions options = VortexOptions.of(Map.of()); - Configuration conf = new Configuration(); - conf.set("fs.file.impl", VortexNestedSchemaPruningIoTest.CountingLocalFileSystem.class.getName()); - conf.set("fs.file.impl.disable.cache", "true"); - VortexPartitionReaderFactory factory = new VortexPartitionReaderFactory( - new FileSourceOptions(scala.collection.immutable.Map$.MODULE$.empty()), - VortexIo.create(options, conf), - options, - SCHEMA, - SCHEMA, - new StructType(), - new Filter[0], - false); - VortexReadMetrics metrics = VortexReadMetrics.forTask(factory); - long rows = 0; - for (String path : files) { - long length = Path.of(path).toFile().length(); - PartitionedFile file = new PartitionedFile( - InternalRow.empty(), - SparkPath.fromPathString("file://" + path), - 0, - length, - new String[0], - 0, - length, - PartitionedFile.apply$default$8()); - try (PartitionReader reader = factory.buildColumnarReader(file, metrics)) { - while (reader.next()) { - rows += reader.get().numRows(); - } - } catch (IOException e) { - throw new UncheckedIOException(e); - } - metrics.publish(); - } - return new long[] {rows, metrics.currentMetricsValues()[0].value()}; + 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 read = + spark.read().format(OneFileAtATimeSource.class.getName()).load(path); + verify(read, 16384); } private Dataset read() { From bc85c7cd9f4533b6aaf3d5e71d19e116a84d4886 Mon Sep 17 00:00:00 2001 From: Christopher Wirt Date: Wed, 11 Nov 2026 11:02:41 +0000 Subject: [PATCH 3/3] docs(spark): document buildColumnarReader(file)'s metrics contract States the task-thread requirement and that the custom metric is one cumulative total per task. Looks up TaskContext once per reader. Co-Authored-By: Claude Opus 5.5 --- .../read/VortexAggregateReaderFactory.java | 7 +++++++ .../read/VortexPartitionReaderFactory.java | 7 +++++++ .../vortex/spark/read/VortexReadMetrics.java | 17 +++++++++-------- 3 files changed, 23 insertions(+), 8 deletions(-) diff --git a/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexAggregateReaderFactory.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexAggregateReaderFactory.java index 2bfa2bd09de..7edc630945b 100644 --- a/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexAggregateReaderFactory.java +++ b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexAggregateReaderFactory.java @@ -81,6 +81,13 @@ public PartitionReader 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 buildColumnarReader(PartitionedFile file) { return VortexReadMetrics.createFileReader(bytesRead -> buildColumnarReader(file, bytesRead)); diff --git a/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexPartitionReaderFactory.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexPartitionReaderFactory.java index 638fe38896e..3a07ab2cec6 100644 --- a/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexPartitionReaderFactory.java +++ b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexPartitionReaderFactory.java @@ -76,6 +76,13 @@ public PartitionReader 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 buildColumnarReader(PartitionedFile file) { return VortexReadMetrics.createFileReader(bytesRead -> buildColumnarReader(file, bytesRead)); diff --git a/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexReadMetrics.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexReadMetrics.java index 70f5c587f56..ad5b46b658e 100644 --- a/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexReadMetrics.java +++ b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexReadMetrics.java @@ -115,8 +115,9 @@ static PartitionReader createReader( Object factory, FileSourceOptions options, BiFunction> buildReader) { + TaskContext context = TaskContext.get(); // Several input partitions can run in one coalesced task, so each reader of this scan shares one total. - VortexReadMetrics metrics = forScan(factory); + VortexReadMetrics metrics = forScan(context, factory); PartitionReader files = new FilePartitionReaderFactory() { @Override public FileSourceOptions options() { @@ -135,7 +136,7 @@ public PartitionReader buildColumnarReader(PartitionedFile file) }.createColumnarReader(partition); // FilePartitionReader does not forward currentMetricsValues from individual files. - return reporting(files, metrics); + return reporting(context, files, metrics); } /** @@ -144,19 +145,19 @@ public PartitionReader buildColumnarReader(PartitionedFile file) */ static PartitionReader createFileReader( Function> buildReader) { - VortexReadMetrics metrics = forScan(SINGLE_FILES); - return reporting(buildReader.apply(metrics), metrics); + 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(Object scan) { - TaskBytes taskBytes = taskBytes(TaskContext.get()); + 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 reporting( - PartitionReader files, VortexReadMetrics metrics) { - TaskContext context = TaskContext.get(); + TaskContext context, PartitionReader files, VortexReadMetrics metrics) { return new PartitionReader<>() { @Override public boolean next() throws IOException {