diff --git a/java/vortex-jni/src/main/java/dev/vortex/io/PooledReadable.java b/java/vortex-jni/src/main/java/dev/vortex/io/PooledReadable.java index e3ab92f0709..8a9e8d61166 100644 --- a/java/vortex-jni/src/main/java/dev/vortex/io/PooledReadable.java +++ b/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. * - *

Streams are opened on demand and reused until this readable is closed. Subclasses supply only - * {@link #openStream()}. + *

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 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 { diff --git a/java/vortex-spark/common/src/main/java/dev/vortex/spark/io/HadoopReadable.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/io/HadoopReadable.java index 5f071004448..44839c82354 100644 --- a/java/vortex-spark/common/src/main/java/dev/vortex/spark/io/HadoopReadable.java +++ b/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 diff --git a/java/vortex-spark/common/src/main/java/dev/vortex/spark/io/VortexIo.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/io/VortexIo.java index e3b82466180..7d867bd2c91 100644 --- a/java/vortex-spark/common/src/main/java/dev/vortex/spark/io/VortexIo.java +++ b/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); } } 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 79aaa7fde5b..ddf0b8bafa5 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 @@ -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 buildReader(PartitionedFile file) { throw new UnsupportedOperationException("row-based aggregate reads are not supported"); } + @Override + public PartitionReader createColumnarReader(InputPartition partition) { + return VortexReadMetrics.createReader(partition, this, fileOptions, this::buildColumnarReader); + } + @Override public PartitionReader buildColumnarReader(PartitionedFile file) { + return buildColumnarReader(file, ignored -> {}); + } + + private PartitionReader 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. " diff --git a/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexBytesReadMetric.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexBytesReadMetric.java new file mode 100644 index 00000000000..7ce28fcdf6c --- /dev/null +++ b/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"; + } +} diff --git a/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexFooterReader.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexFooterReader.java index 8ac3e4c8301..0b3cc3d1530 100644 --- a/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexFooterReader.java +++ b/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 List mapFooters( private static T withDataSource( VortexFile file, VortexIo io, VortexOptions options, DataSourceFunction function) { + return withDataSource(file, io, options, ignored -> {}, function); + } + + private static T withDataSource( + VortexFile file, + VortexIo io, + VortexOptions options, + LongConsumer bytesRead, + DataSourceFunction 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) { diff --git a/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexPartitionReader.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexPartitionReader.java index 2fa75f9f965..2fd70f71202 100644 --- a/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexPartitionReader.java +++ b/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 arrowFields = dataSource.arrowSchema(allocator).getFields(); 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 13f8ee06922..5091c878869 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 @@ -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 buildReader(PartitionedFile file) { throw new UnsupportedOperationException("row-based V2 reads are not supported"); } + @Override + public PartitionReader createColumnarReader(InputPartition partition) { + return VortexReadMetrics.createReader(partition, this, fileOptions, this::buildColumnarReader); + } + @Override public PartitionReader buildColumnarReader(PartitionedFile file) { + return buildColumnarReader(file, ignored -> {}); + } + + private PartitionReader 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 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 new file mode 100644 index 00000000000..7ae5f08a2cf --- /dev/null +++ b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexReadMetrics.java @@ -0,0 +1,178 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +package dev.vortex.spark.read; + +import java.io.IOException; +import java.io.UncheckedIOException; +import java.util.IdentityHashMap; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import java.util.concurrent.atomic.AtomicLong; +import java.util.function.BiFunction; +import java.util.function.LongConsumer; +import org.apache.hadoop.fs.FileSystem; +import org.apache.spark.TaskContext; +import org.apache.spark.sql.catalyst.FileSourceOptions; +import org.apache.spark.sql.catalyst.InternalRow; +import org.apache.spark.sql.connector.metric.CustomTaskMetric; +import org.apache.spark.sql.connector.read.InputPartition; +import org.apache.spark.sql.connector.read.PartitionReader; +import org.apache.spark.sql.execution.datasources.PartitionedFile; +import org.apache.spark.sql.execution.datasources.v2.FilePartitionReaderFactory; +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 { + 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") + private static final FileSystem.Statistics NATIVE_STATISTICS = + FileSystem.getStatistics("vortex-native", NativeReadStatistics.class); + + private abstract static class NativeReadStatistics extends FileSystem {} + + private final AtomicLong bytesRead; + private final TaskBytes taskBytes; + private final Thread taskThread = Thread.currentThread(); + + VortexReadMetrics() { + this(new AtomicLong(), new TaskBytes()); + } + + private VortexReadMetrics(AtomicLong bytesRead, TaskBytes taskBytes) { + this.bytesRead = bytesRead; + this.taskBytes = taskBytes; + } + + @Override + public void accept(long bytes) { + bytesRead.addAndGet(bytes); + if (Thread.currentThread() != taskThread) { + taskBytes.workerBytes.addAndGet(bytes); + } + } + + private static TaskBytes taskBytes(TaskContext context) { + if (context == null) { + return new TaskBytes(); + } + TaskBytes created = new TaskBytes(); + TaskBytes existing = TASK_BYTES.putIfAbsent(context, created); + if (existing != null) { + return existing; + } + // Registered outside the map update: Spark runs this immediately if the task has already completed. + context.addTaskCompletionListener((TaskCompletionListener) ignored -> TASK_BYTES.remove(context)); + return created; + } + + private void publishWorkerBytes(TaskContext context) { + if (context == null) { + return; + } + long completed = taskBytes.workerBytes.get(); + long delta = completed - taskBytes.publishedWorkerBytes; + if (delta != 0) { + // Spark overwrites input bytes from the task thread's Hadoop statistics, so worker-thread bytes must + // land there too. Task-thread reads are already counted by the storage FileSystem's own statistics. + NATIVE_STATISTICS.incrementBytesRead(delta); + context.taskMetrics().inputMetrics().incBytesRead(delta); + taskBytes.publishedWorkerBytes = completed; + } + } + + private static final class TaskBytes { + private final AtomicLong workerBytes = new AtomicLong(); + private final IdentityHashMap scanBytes = new IdentityHashMap<>(); + // Accessed only on the Spark task thread, shared by every reader/scan consumed by that task. + private long publishedWorkerBytes; + } + + CustomTaskMetric[] currentMetricsValues() { + long snapshot = bytesRead.get(); + return new CustomTaskMetric[] { + new CustomTaskMetric() { + @Override + public String name() { + return VortexBytesReadMetric.NAME; + } + + @Override + public long value() { + return snapshot; + } + } + }; + } + + static PartitionReader createReader( + InputPartition partition, + Object factory, + 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); + PartitionReader files = new FilePartitionReaderFactory() { + @Override + public FileSourceOptions options() { + return options; + } + + @Override + public PartitionReader buildReader(PartitionedFile file) { + throw new UnsupportedOperationException("row-based V2 reads are not supported"); + } + + @Override + public PartitionReader buildColumnarReader(PartitionedFile file) { + return buildReader.apply(file, metrics); + } + }.createColumnarReader(partition); + + // FilePartitionReader does not forward currentMetricsValues from individual files. + return new PartitionReader<>() { + @Override + public boolean next() throws IOException { + try { + return files.next(); + } finally { + metrics.publishWorkerBytes(context); + } + } + + @Override + public ColumnarBatch get() { + return files.get(); + } + + @Override + public CustomTaskMetric[] currentMetricsValues() { + if (context != null && context.isCompleted()) { + // Spark asks for its final custom snapshot before closing the reader. Drain read-ahead first. + try { + files.close(); + } catch (IOException e) { + throw new UncheckedIOException(e); + } + } + metrics.publishWorkerBytes(context); + return metrics.currentMetricsValues(); + } + + @Override + public void close() throws IOException { + try { + files.close(); + } finally { + metrics.publishWorkerBytes(context); + } + } + }; + } +} diff --git a/java/vortex-spark/common/src/main/scala/dev/vortex/spark/read/VortexScan.scala b/java/vortex-spark/common/src/main/scala/dev/vortex/spark/read/VortexScan.scala index 71e35da5df2..2a54df49300 100644 --- a/java/vortex-spark/common/src/main/scala/dev/vortex/spark/read/VortexScan.scala +++ b/java/vortex-spark/common/src/main/scala/dev/vortex/spark/read/VortexScan.scala @@ -14,6 +14,7 @@ import org.apache.spark.sql.catalyst.FileSourceOptions import org.apache.spark.sql.catalyst.expressions.Expression import org.apache.spark.sql.connector.expressions.aggregate.Aggregation import org.apache.spark.sql.connector.read.{PartitionReaderFactory, Statistics} +import org.apache.spark.sql.connector.metric.CustomMetric import org.apache.spark.sql.execution.datasources.{ AggregatePushDownUtils, PartitioningAwareFileIndex @@ -101,6 +102,9 @@ final class VortexScan( } } + override def supportedCustomMetrics(): Array[CustomMetric] = + Array(new VortexBytesReadMetric()) + override def estimateStatistics(): Statistics = cachedStatistics private lazy val cachedStatistics: Statistics = { 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 new file mode 100644 index 00000000000..4a65a3b6f24 --- /dev/null +++ b/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexInputMetricsTest.java @@ -0,0 +1,143 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +package dev.vortex.spark; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.nio.file.Path; +import java.util.concurrent.atomic.AtomicLong; +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.execution.SparkPlan; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.StructType; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestInstance; +import org.junit.jupiter.api.io.TempDir; + +/** Checks the final task metrics Spark publishes, after its own Hadoop metric update. */ +@TestInstance(TestInstance.Lifecycle.PER_CLASS) +final class VortexInputMetricsTest { + private static final StructType SCHEMA = new StructType().add("id", DataTypes.LongType); + private SparkSession spark; + private String table; + private final InputBytesListener listener = new InputBytesListener(); + + @BeforeAll + void setUp(@TempDir Path tempDir) { + spark = SparkSession.builder() + .appName("VortexInputMetricsTest") + .master("local[2]") + .config("spark.driver.host", "127.0.0.1") + .config("spark.ui.enabled", "false") + .config("spark.sql.adaptive.enabled", "false") + .config("spark.sql.sources.useV1SourceList", "") + .config("spark.sql.files.openCostInBytes", "0") + .config("spark.sql.files.minPartitionNum", "1") + .config( + "spark.hadoop.fs.file.impl", + VortexNestedSchemaPruningIoTest.CountingLocalFileSystem.class.getName()) + .config("spark.hadoop.fs.file.impl.disable.cache", "true") + .getOrCreate(); + table = tempDir.resolve("table").toString(); + spark.range(0, 16384, 1, 4).write().format("vortex").save(table); + spark.sparkContext().addSparkListener(listener); + } + + @AfterAll + void tearDown() { + if (spark != null) { + spark.stop(); + } + } + + @Test + void multipleFilesAccumulateWithoutDoubleCounting() throws Exception { + Dataset read = read(); + assertEquals(1, read.rdd().getNumPartitions(), "several files must share one partition"); + verify(read, 16384); + } + + @Test + void coalescedTaskAccumulatesSeveralInputPartitions() throws Exception { + spark.conf().set("spark.sql.files.maxPartitionBytes", "1024"); + try { + Dataset read = read(); + assertTrue(read.rdd().getNumPartitions() > 1); + verify(read.coalesce(1), 16384); + } finally { + spark.conf().unset("spark.sql.files.maxPartitionBytes"); + } + } + + @Test + void separateScansKeepIndependentCustomTotalsInOneTask() throws Exception { + verify(read().filter("id % 3 = 0").union(read().filter("id % 3 != 0")).coalesce(1), 16384); + } + + @Test + void emptyResultsStillReportFooterAndFilterReads() throws Exception { + verify(read().filter("id % 2 = -1"), 0); + } + + @Test + void pushedCountReportsFooterBytes() throws Exception { + Dataset count = read().selectExpr("count(*)"); + assertTrue(count.queryExecution().executedPlan().toString().contains("PushedAggregation: [COUNT(*)]")); + verify(count, 1); + } + + @Test + void earlyTerminationReportsInputBytes() throws Exception { + verify(read().limit(1), 1); + } + + private Dataset read() { + return spark.read() + .schema(SCHEMA) + .format("vortex") + .option("vortex.stats.rowCount", "false") + .load(table); + } + + private void verify(Dataset query, int rows) throws Exception { + SparkPlan plan = query.queryExecution().executedPlan(); + assertTrue(plan.toString().contains("BatchScan"), plan.toString()); + spark.sparkContext().listenerBus().waitUntilEmpty(10000); + listener.bytes.set(0); + VortexNestedSchemaPruningIoTest.BYTES_READ.set(0); + assertEquals(rows, query.collectAsList().size()); + spark.sparkContext().listenerBus().waitUntilEmpty(10000); + long source = VortexNestedSchemaPruningIoTest.BYTES_READ.get(); + assertTrue(source > 0, "the scan must read source bytes"); + assertEquals(source, listener.bytes.get(), "task metrics must include each source byte once"); + assertEquals(source, customBytes(plan), "custom metrics must retain bytes from every file"); + } + + private static long customBytes(SparkPlan plan) { + long bytes = plan.metrics().contains("vortexBytesRead") + ? plan.metrics().apply("vortexBytesRead").value() + : 0; + var children = plan.children().iterator(); + while (children.hasNext()) { + bytes += customBytes(children.next()); + } + return bytes; + } + + private static final class InputBytesListener extends SparkListener { + private final AtomicLong bytes = new AtomicLong(); + + @Override + public void onTaskEnd(SparkListenerTaskEnd event) { + bytes.addAndGet(event.taskMetrics().inputMetrics().bytesRead()); + } + } +} diff --git a/java/vortex-spark/common/src/test/java/dev/vortex/spark/io/HadoopReadableTest.java b/java/vortex-spark/common/src/test/java/dev/vortex/spark/io/HadoopReadableTest.java index e44ea9506fe..5bd76a0f21e 100644 --- a/java/vortex-spark/common/src/test/java/dev/vortex/spark/io/HadoopReadableTest.java +++ b/java/vortex-spark/common/src/test/java/dev/vortex/spark/io/HadoopReadableTest.java @@ -25,14 +25,21 @@ import java.util.UUID; import java.util.concurrent.Callable; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicLong; +import java.util.function.LongConsumer; import org.apache.hadoop.conf.Configuration; import org.apache.hadoop.fs.FSDataInputStream; +import org.apache.hadoop.fs.FSInputStream; import org.apache.hadoop.fs.FileStatus; import org.apache.hadoop.fs.FileSystem; import org.apache.hadoop.fs.FutureDataInputStreamBuilder; +import org.apache.hadoop.fs.LocalFileSystem; import org.apache.hadoop.fs.RawLocalFileSystem; import org.apache.hadoop.fs.impl.FutureDataInputStreamBuilderImpl; import org.junit.jupiter.api.Test; @@ -149,8 +156,9 @@ void concurrentReadsEachGetTheirOwnStream() throws Exception { Path file = write("data.vortex", content); int readers = 8; int chunk = SIZE / readers; + AtomicLong bytes = new AtomicLong(); - try (NativeReadable readable = HadoopReadable.open(conf, file.toString())) { + try (NativeReadable readable = openCounting(conf, file.toString(), SIZE, bytes::addAndGet)) { ExecutorService pool = Executors.newFixedThreadPool(readers); try { List> reads = new ArrayList<>(); @@ -177,6 +185,7 @@ void concurrentReadsEachGetTheirOwnStream() throws Exception { } finally { pool.shutdownNow(); } + assertEquals((long) SIZE * 20, bytes.get()); } } @@ -190,6 +199,147 @@ void readsPastTheEndFail() throws IOException { } } + @Test + void failedReadsDoNotReportRequestedBytes() throws IOException { + Path file = write("short.vortex", content(16)); + AtomicLong bytes = new AtomicLong(); + // The trusted length exceeds the storage length, so this reaches the stream's EOF guard. + try (NativeReadable readable = openCounting(conf, file.toString(), 32, bytes::addAndGet)) { + assertThrows(EOFException.class, () -> readable.readFully(0, ByteBuffer.allocateDirect(32))); + assertEquals(0, bytes.get()); + readable.readFully(0, ByteBuffer.allocateDirect(16)); + assertEquals(16, bytes.get()); + } + } + + @Test + void closeDrainsInFlightReads() throws Exception { + closeDuringRead(false); + } + + @Test + void interruptedCloseDoesNotWaitForAStalledRead() throws Exception { + closeDuringRead(true); + } + + @Test + void interruptedCloseSucceedsWithNoReadsInFlight() throws IOException { + Path file = write("idle.vortex", content(16)); + NativeReadable readable = openCounting(conf, file.toString(), 16, ignored -> {}); + readable.readFully(0, ByteBuffer.allocateDirect(16)); + Thread.currentThread().interrupt(); + try { + readable.close(); + assertTrue(Thread.currentThread().isInterrupted()); + } finally { + Thread.interrupted(); + } + } + + private void closeDuringRead(boolean interrupt) throws Exception { + BlockingLocalFileSystem.reading = new CountDownLatch(1); + BlockingLocalFileSystem.finish = new CountDownLatch(1); + Configuration blockingConf = new Configuration(conf); + blockingConf.setClass("fs.file.impl", BlockingLocalFileSystem.class, org.apache.hadoop.fs.FileSystem.class); + blockingConf.setBoolean("fs.file.impl.disable.cache", true); + Path file = write("blocked.vortex", content(16)); + AtomicLong bytes = new AtomicLong(); + NativeReadable readable = openCounting(blockingConf, file.toString(), 16, bytes::addAndGet); + ExecutorService workers = Executors.newFixedThreadPool(2); + try { + Future read = workers.submit(() -> { + readable.readFully(0, ByteBuffer.allocateDirect(16)); + return null; + }); + assertTrue(BlockingLocalFileSystem.reading.await(10, TimeUnit.SECONDS)); + Future close = workers.submit(() -> { + if (interrupt) { + Thread.currentThread().interrupt(); + try { + assertThrows(InterruptedIOException.class, readable::close); + assertTrue(Thread.currentThread().isInterrupted()); + } finally { + Thread.interrupted(); + } + } else { + readable.close(); + } + return null; + }); + if (interrupt) { + close.get(10, TimeUnit.SECONDS); + } else { + assertThrows(TimeoutException.class, () -> close.get(100, TimeUnit.MILLISECONDS)); + } + BlockingLocalFileSystem.finish.countDown(); + read.get(10, TimeUnit.SECONDS); + close.get(10, TimeUnit.SECONDS); + assertEquals(16, bytes.get()); + assertThrows(IllegalStateException.class, () -> readable.readFully(0, ByteBuffer.allocateDirect(16))); + } finally { + BlockingLocalFileSystem.finish.countDown(); + workers.shutdownNow(); + readable.close(); + } + } + + public static final class BlockingLocalFileSystem extends LocalFileSystem { + private static CountDownLatch reading; + private static CountDownLatch finish; + + @Override + public FSDataInputStream open(org.apache.hadoop.fs.Path path, int bufferSize) throws IOException { + return new FSDataInputStream(new BlockingInputStream(super.open(path, bufferSize))); + } + } + + private static final class BlockingInputStream extends FSInputStream { + private final FSDataInputStream input; + + private BlockingInputStream(FSDataInputStream input) { + this.input = input; + } + + @Override + public void seek(long position) throws IOException { + input.seek(position); + } + + @Override + public long getPos() throws IOException { + return input.getPos(); + } + + @Override + public boolean seekToNewSource(long target) throws IOException { + return input.seekToNewSource(target); + } + + @Override + public int read() throws IOException { + return input.read(); + } + + @Override + public int read(byte[] buffer, int offset, int length) throws IOException { + BlockingLocalFileSystem.reading.countDown(); + try { + if (!BlockingLocalFileSystem.finish.await(10, TimeUnit.SECONDS)) { + throw new IOException("read timed out"); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IOException(e); + } + return input.read(buffer, offset, length); + } + + @Override + public void close() throws IOException { + input.close(); + } + } + @Test void readingAfterCloseFails() throws IOException { Path file = write("data.vortex", content(1024)); @@ -282,6 +432,16 @@ public CompletableFuture build() throws IOException { } } + private static NativeReadable openCounting(Configuration conf, String path, long length, LongConsumer bytesRead) { + return HadoopReadable.open( + conf, + path, + length, + NativeReadable.DEFAULT_COALESCE_DISTANCE, + NativeReadable.DEFAULT_COALESCE_MAX_SIZE, + bytesRead); + } + private Path write(String name, byte[] content) throws IOException { Path file = tempDir.resolve(name); Files.write(file, content); diff --git a/java/vortex-spark/common/src/test/java/dev/vortex/spark/read/VortexReadMetricsTest.java b/java/vortex-spark/common/src/test/java/dev/vortex/spark/read/VortexReadMetricsTest.java new file mode 100644 index 00000000000..650bc0b23c7 --- /dev/null +++ b/java/vortex-spark/common/src/test/java/dev/vortex/spark/read/VortexReadMetricsTest.java @@ -0,0 +1,127 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +package dev.vortex.spark.read; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; +import org.apache.spark.TaskContext$; +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.read.PartitionReader; +import org.apache.spark.sql.execution.datasources.FilePartition; +import org.apache.spark.sql.execution.datasources.PartitionedFile; +import org.apache.spark.sql.vectorized.ColumnarBatch; +import org.apache.spark.util.TaskCompletionListener; +import org.junit.jupiter.api.Test; +import scala.Option; +import scala.collection.immutable.Map$; + +final class VortexReadMetricsTest { + @Test + void completionSnapshotIncludesReadsThatFinishDuringClose() throws Exception { + var context = TaskContext$.MODULE$.empty(); + TaskContext$.MODULE$.setTaskContext(context); + ExecutorService worker = Executors.newSingleThreadExecutor(); + CountDownLatch reading = new CountDownLatch(1); + CountDownLatch finish = new CountDownLatch(1); + AtomicReference> read = new AtomicReference<>(); + try { + PartitionedFile file = new PartitionedFile( + InternalRow.empty(), + SparkPath.fromPathString("file:///memory.vortex"), + 0, + 23, + new String[0], + 0, + 23, + PartitionedFile.apply$default$8()); + try (PartitionReader reader = VortexReadMetrics.createReader( + new FilePartition(0, new PartitionedFile[] {file}), + new Object(), + new FileSourceOptions(Map$.MODULE$.empty()), + (ignored, metrics) -> new PartitionReader<>() { + private boolean closed; + + @Override + public boolean next() { + read.set(worker.submit(() -> { + reading.countDown(); + if (!finish.await(10, TimeUnit.SECONDS)) { + throw new IOException("read timed out"); + } + metrics.accept(23); + return null; + })); + return true; + } + + @Override + public ColumnarBatch get() { + throw new UnsupportedOperationException(); + } + + @Override + public void close() throws IOException { + if (!closed) { + closed = true; + finish.countDown(); + try { + read.get().get(10, TimeUnit.SECONDS); + } catch (Exception e) { + throw new IOException(e); + } + } + } + })) { + assertTrue(reader.next()); + assertTrue(reading.await(10, TimeUnit.SECONDS)); + assertEquals(0, reader.currentMetricsValues()[0].value()); + context.addTaskCompletionListener((TaskCompletionListener) ignored -> { + assertEquals(23, reader.currentMetricsValues()[0].value()); + assertEquals(23, reader.currentMetricsValues()[0].value()); + }); + context.markTaskCompleted(Option.empty()); + assertEquals(23, context.taskMetrics().inputMetrics().bytesRead()); + } + } finally { + finish.countDown(); + worker.shutdownNow(); + TaskContext$.MODULE$.unset(); + } + } + + @Test + void cumulativeSnapshotsIncludeConcurrentReads() throws Exception { + VortexReadMetrics metrics = new VortexReadMetrics(); + ExecutorService workers = Executors.newFixedThreadPool(4); + try { + List> reads = new ArrayList<>(); + for (int i = 0; i < 100; i++) { + reads.add(workers.submit(() -> metrics.accept(17))); + } + for (Future read : reads) { + read.get(); + } + } finally { + workers.shutdownNow(); + } + var snapshot = metrics.currentMetricsValues()[0]; + assertEquals(1700, snapshot.value()); + assertEquals("vortexBytesRead", snapshot.name()); + metrics.accept(3); + assertEquals(1700, snapshot.value()); + assertEquals(1703, metrics.currentMetricsValues()[0].value()); + } +}