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..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,9 +81,16 @@ 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 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 5091c878869..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,9 +76,16 @@ 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 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/VortexReadMetrics.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexReadMetrics.java index 7ae5f08a2cf..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 @@ -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 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 createReader( FileSourceOptions options, 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 shares one total. + VortexReadMetrics metrics = forScan(context, factory); PartitionReader files = new FilePartitionReaderFactory() { @Override public FileSourceOptions options() { @@ -136,6 +136,28 @@ public PartitionReader 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 createFileReader( + Function> 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 reporting( + TaskContext context, PartitionReader files, VortexReadMetrics metrics) { return new PartitionReader<>() { @Override public boolean next() throws IOException { 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 4a65a3b6f24..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 @@ -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 read = + spark.read().format(OneFileAtATimeSource.class.getName()).load(path); + verify(read, 16384); + } + private Dataset read() { return spark.read() .schema(SCHEMA)