diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexCountStarBatchExec.java b/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexCountStarBatchExec.java new file mode 100644 index 00000000000..b24bda1f40e --- /dev/null +++ b/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexCountStarBatchExec.java @@ -0,0 +1,47 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +package dev.vortex.spark.read; + +import dev.vortex.spark.VortexSparkSession; +import java.util.List; +import java.util.Map; +import org.apache.spark.sql.connector.read.Batch; +import org.apache.spark.sql.connector.read.InputPartition; +import org.apache.spark.sql.connector.read.PartitionReaderFactory; + +/** + * Spark V2 {@link Batch} for a pushed-down {@code COUNT(*)} over Vortex files. + * + *

Plans one {@link VortexCountStarInputPartition} per resolved file so the per-file footer reads run in parallel on + * executors, mirroring the per-file partitioning of {@link VortexBatchExec}. Paths are resolved with the same directory + * expansion as regular scans so a count and a scan over the same relation always see the same file set. + */ +public final class VortexCountStarBatchExec implements Batch { + private final List paths; + private final Map formatOptions; + + /** + * Creates a new VortexCountStarBatchExec for the specified file paths. The caller is responsible for passing + * immutable collections; the constructor does not copy. + * + * @param paths the list of Vortex file paths to count over + * @param formatOptions the format options for opening the files + */ + VortexCountStarBatchExec(List paths, Map formatOptions) { + this.paths = paths; + this.formatOptions = formatOptions; + } + + @Override + public InputPartition[] planInputPartitions() { + List resolvedPaths = + VortexBatchExec.resolveVortexPaths(VortexSparkSession.get(formatOptions), paths, formatOptions); + return resolvedPaths.stream().map(VortexCountStarInputPartition::new).toArray(InputPartition[]::new); + } + + @Override + public PartitionReaderFactory createReaderFactory() { + return new VortexCountStarPartitionReaderFactory(formatOptions); + } +} diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexCountStarInputPartition.java b/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexCountStarInputPartition.java new file mode 100644 index 00000000000..2112fa54c89 --- /dev/null +++ b/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexCountStarInputPartition.java @@ -0,0 +1,18 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +package dev.vortex.spark.read; + +import org.apache.spark.sql.connector.read.InputPartition; + +/** + * Input partition for a pushed-down {@code COUNT(*)}: a single Vortex file whose footer row count is read on an + * executor. + * + *

Carrying exactly one file per partition is what makes the count exact: a single-file + * {@link dev.vortex.api.DataSource} opens its only file eagerly, so its row count comes straight from the footer rather + * than the multi-file extrapolation documented on {@link VortexScan#estimateStatistics()}. + * + * @param path the resolved {@code .vortex} file path + */ +public record VortexCountStarInputPartition(String path) implements InputPartition {} diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexCountStarPartitionReader.java b/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexCountStarPartitionReader.java new file mode 100644 index 00000000000..114dd27042c --- /dev/null +++ b/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexCountStarPartitionReader.java @@ -0,0 +1,79 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +package dev.vortex.spark.read; + +import static com.google.common.base.Preconditions.checkState; + +import dev.vortex.api.DataSource; +import dev.vortex.spark.VortexSparkSession; +import java.util.List; +import java.util.Map; +import org.apache.spark.sql.connector.read.PartitionReader; +import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector; +import org.apache.spark.sql.vectorized.ColumnarBatch; + +/** + * Reads the exact row count of a single Vortex file from its footer and emits it as a one-row {@link ColumnarBatch}. + * + *

No column data is decoded: opening a single-file {@link DataSource} reads the file footer eagerly, so + * {@link DataSource#rowCount()} is {@link DataSource.RowCount.Exact}. The reader fails fast if the count is ever not + * exact rather than silently returning a wrong result — a single-file source is exact by construction today, and this + * guard turns any future change of that invariant into a loud error instead of a correctness bug. + */ +public final class VortexCountStarPartitionReader implements PartitionReader { + private final String path; + private final Map formatOptions; + + private ColumnarBatch batch; + private boolean emitted = false; + + /** + * Creates a new reader for a single file. + * + * @param path the resolved {@code .vortex} file path + * @param formatOptions the format options for opening the file + */ + VortexCountStarPartitionReader(String path, Map formatOptions) { + this.path = path; + this.formatOptions = formatOptions; + } + + @Override + public boolean next() { + if (emitted) { + return false; + } + long count = readFooterRowCount(); + OnHeapColumnVector[] vectors = OnHeapColumnVector.allocateColumns(1, VortexCountStarScan.COUNT_SCHEMA); + vectors[0].putLong(0, count); + this.batch = new ColumnarBatch(vectors, 1); + this.emitted = true; + return true; + } + + @Override + public ColumnarBatch get() { + checkState(batch != null, "next() must return true before get()"); + return batch; + } + + @Override + public void close() { + if (batch != null) { + batch.close(); + batch = null; + } + } + + private long readFooterRowCount() { + DataSource source = DataSource.open(VortexSparkSession.get(formatOptions), List.of(path), formatOptions); + DataSource.RowCount rowCount = source.rowCount(); + checkState( + rowCount instanceof DataSource.RowCount.Exact, + "expected exact footer row count for single-file data source %s, got %s", + path, + rowCount); + return ((DataSource.RowCount.Exact) rowCount).value(); + } +} diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexCountStarPartitionReaderFactory.java b/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexCountStarPartitionReaderFactory.java new file mode 100644 index 00000000000..ab8e9b68aab --- /dev/null +++ b/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexCountStarPartitionReaderFactory.java @@ -0,0 +1,51 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +package dev.vortex.spark.read; + +import static com.google.common.base.Preconditions.checkArgument; + +import java.util.Map; +import org.apache.spark.sql.catalyst.InternalRow; +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.vectorized.ColumnarBatch; + +/** + * Factory for {@link VortexCountStarPartitionReader}s. + * + *

Columnar-only, mirroring {@link VortexPartitionReaderFactory}: the count is emitted as a single one-row + * {@link ColumnarBatch}. + */ +public final class VortexCountStarPartitionReaderFactory implements PartitionReaderFactory { + private final Map formatOptions; + + /** + * Creates a new factory. + * + * @param formatOptions the format options for opening the files + */ + VortexCountStarPartitionReaderFactory(Map formatOptions) { + this.formatOptions = formatOptions; + } + + @Override + public PartitionReader createReader(InputPartition partition) { + throw new UnsupportedOperationException("row-based reads are not supported"); + } + + @Override + public PartitionReader createColumnarReader(InputPartition partition) { + checkArgument( + partition instanceof VortexCountStarInputPartition, + "expected VortexCountStarInputPartition, got %s", + partition.getClass().getName()); + return new VortexCountStarPartitionReader(((VortexCountStarInputPartition) partition).path(), formatOptions); + } + + @Override + public boolean supportColumnarReads(InputPartition partition) { + return true; + } +} diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexCountStarScan.java b/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexCountStarScan.java new file mode 100644 index 00000000000..1226f7d7691 --- /dev/null +++ b/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexCountStarScan.java @@ -0,0 +1,59 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +package dev.vortex.spark.read; + +import java.util.List; +import java.util.Map; +import org.apache.spark.sql.connector.read.Batch; +import org.apache.spark.sql.connector.read.Scan; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.StructType; + +/** + * Spark V2 {@link Scan} that answers a pushed-down global {@code COUNT(*)} from Vortex file footer metadata. + * + *

Built by {@link VortexScanBuilder} when Spark pushes an aggregation that is exactly one {@code COUNT(*)} with no + * grouping expressions and no pushed predicates. Instead of scanning data, each input partition opens a single-file + * {@link dev.vortex.api.DataSource} and reads the exact row count recorded in that file's footer, emitting one partial + * count per file. This is a partial pushdown: Spark performs the final summation of the per-file counts, following the + * column-order contract of {@link org.apache.spark.sql.connector.read.SupportsPushDownAggregates}. + */ +public final class VortexCountStarScan implements Scan { + static final StructType COUNT_SCHEMA = new StructType().add("count(*)", DataTypes.LongType, false); + + private final List paths; + private final Map formatOptions; + + /** + * Creates a new VortexCountStarScan for the specified file paths. The caller is responsible for passing immutable + * collections; the constructor does not copy. + * + * @param paths the list of Vortex file paths to count over + * @param formatOptions the format options for opening the files + */ + VortexCountStarScan(List paths, Map formatOptions) { + this.paths = paths; + this.formatOptions = formatOptions; + } + + @Override + public StructType readSchema() { + return COUNT_SCHEMA; + } + + @Override + public Batch toBatch() { + return new VortexCountStarBatchExec(paths, formatOptions); + } + + @Override + public String description() { + return "VortexCountStarScan PushedAggregation: [COUNT(*)]"; + } + + @Override + public ColumnarSupportMode columnarSupportMode() { + return ColumnarSupportMode.SUPPORTED; + } +} diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexScanBuilder.java b/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexScanBuilder.java index 62c8085aa0f..8c653d87262 100644 --- a/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexScanBuilder.java +++ b/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexScanBuilder.java @@ -19,9 +19,13 @@ import org.apache.spark.sql.connector.catalog.Column; import org.apache.spark.sql.connector.expressions.NamedReference; import org.apache.spark.sql.connector.expressions.Transform; +import org.apache.spark.sql.connector.expressions.aggregate.AggregateFunc; +import org.apache.spark.sql.connector.expressions.aggregate.Aggregation; +import org.apache.spark.sql.connector.expressions.aggregate.CountStar; import org.apache.spark.sql.connector.expressions.filter.Predicate; import org.apache.spark.sql.connector.read.Scan; import org.apache.spark.sql.connector.read.ScanBuilder; +import org.apache.spark.sql.connector.read.SupportsPushDownAggregates; import org.apache.spark.sql.connector.read.SupportsPushDownRequiredColumns; import org.apache.spark.sql.connector.read.SupportsPushDownV2Filters; import org.apache.spark.sql.types.DataType; @@ -29,13 +33,14 @@ /** Spark V2 {@link ScanBuilder} for table scans over Vortex files. */ public final class VortexScanBuilder - implements ScanBuilder, SupportsPushDownRequiredColumns, SupportsPushDownV2Filters { + implements ScanBuilder, SupportsPushDownRequiredColumns, SupportsPushDownV2Filters, SupportsPushDownAggregates { private final ImmutableList.Builder paths; private final List tableColumns; private final List readColumns; private final Map formatOptions; private final Set partitionColumnNames; private Predicate[] pushedPredicates = new Predicate[0]; + private boolean pushedCountStar = false; /** Creates a new VortexScanBuilder with empty paths and columns. */ public VortexScanBuilder(Map formatOptions) { @@ -120,6 +125,10 @@ public Scan build() { checkState(!paths.isEmpty(), "paths cannot be empty"); + if (pushedCountStar) { + return new VortexCountStarScan(paths, this.formatOptions); + } + return new VortexScan( paths, List.copyOf(this.tableColumns), @@ -128,16 +137,52 @@ public Scan build() { this.formatOptions); } + /** + * Pushes down a global {@code COUNT(*)} so it is answered from file footer metadata instead of scanning data. + * + *

The pushdown is accepted only for an aggregation that is exactly one {@link CountStar} with no grouping + * expressions, and only when no predicates were pushed down (a pushed filter changes the number of matching rows, + * so footer row counts would over-count). Partition-column filters are never pushed (they are evaluated against + * directory paths before planning), so they do not block the pushdown. + * + *

This is a partial pushdown ({@link #supportCompletePushDown(Aggregation)} stays {@code false}): the scan emits + * one partial count per file and Spark performs the final summation. + * + * @param aggregation the aggregation Spark wants to push down + * @return true if the aggregation was pushed, false to fall back to a regular scan + */ + @Override + public boolean pushAggregation(Aggregation aggregation) { + if (pushedPredicates.length != 0) { + return false; + } + if (aggregation.groupByExpressions().length != 0) { + return false; + } + AggregateFunc[] aggregateFunctions = aggregation.aggregateExpressions(); + if (aggregateFunctions.length != 1 || !(aggregateFunctions[0] instanceof CountStar)) { + return false; + } + this.pushedCountStar = true; + return true; + } + /** * Prunes the columns to only include those specified in the required schema. * *

This method clears the current column list and replaces it with columns derived from the required schema. * Currently only supports top-level schema pruning - deeply nested schema pruning is not yet implemented. * + *

After a successful {@link #pushAggregation(Aggregation)} this method is a no-op: Spark re-prunes to the + * aggregation output schema, whose fields do not correspond to any table column. + * * @param requiredSchema the schema specifying which columns are required */ @Override public void pruneColumns(StructType requiredSchema) { + if (pushedCountStar) { + return; + } readColumns.clear(); readColumns.addAll(Arrays.asList(CatalogV2Util.structTypeToV2Columns(requiredSchema))); } diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/VortexCountStarPushDownTest.java b/java/vortex-spark/src/test/java/dev/vortex/spark/VortexCountStarPushDownTest.java new file mode 100644 index 00000000000..cbb4866dad0 --- /dev/null +++ b/java/vortex-spark/src/test/java/dev/vortex/spark/VortexCountStarPushDownTest.java @@ -0,0 +1,198 @@ +// 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.assertFalse; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import dev.vortex.spark.read.VortexCountStarScan; +import dev.vortex.spark.read.VortexScan; +import dev.vortex.spark.read.VortexScanBuilder; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Comparator; +import java.util.Map; +import java.util.stream.Stream; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SaveMode; +import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.connector.expressions.Expression; +import org.apache.spark.sql.connector.expressions.Expressions; +import org.apache.spark.sql.connector.expressions.aggregate.AggregateFunc; +import org.apache.spark.sql.connector.expressions.aggregate.Aggregation; +import org.apache.spark.sql.connector.expressions.aggregate.Count; +import org.apache.spark.sql.connector.expressions.aggregate.CountStar; +import org.apache.spark.sql.connector.expressions.aggregate.Min; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestInstance; +import org.junit.jupiter.api.io.TempDir; + +/** + * Characterizes {@code COUNT(*)} aggregate pushdown: a global count is answered from file footer metadata via + * {@link VortexCountStarScan}, while filtered, grouped, or column-counting aggregations fall back to a regular + * {@link VortexScan} and stay correct. + */ +@TestInstance(TestInstance.Lifecycle.PER_CLASS) +public final class VortexCountStarPushDownTest { + private SparkSession spark; + + @TempDir + Path tempDir; + + @BeforeAll + public void setUp() { + spark = SparkSession.builder() + .appName("VortexCountStarPushDownTest") + .master("local[2]") + .config("spark.driver.host", "127.0.0.1") + .config("spark.sql.shuffle.partitions", "2") + .config("spark.sql.adaptive.enabled", "false") + .config("spark.ui.enabled", "false") + .getOrCreate(); + } + + @AfterAll + public void tearDown() { + if (spark != null) { + spark.stop(); + } + } + + @AfterEach + public void cleanupTempFiles() throws IOException { + if (Files.exists(tempDir)) { + try (Stream walk = Files.walk(tempDir)) { + walk.sorted(Comparator.reverseOrder()).forEach(path -> { + if (!path.equals(tempDir)) { + try { + Files.deleteIfExists(path); + } catch (IOException e) { + // best-effort cleanup + } + } + }); + } + } + } + + @Test + @DisplayName("Global COUNT(*) over a multi-file dataset is answered by VortexCountStarScan and is exact") + public void testCountStarIsPushedDownAndExact() throws IOException { + Path outputPath = writeRows(10_000, "count_multi_file", 4); + + Dataset df = spark.read().format("vortex").load(outputPath.toUri().toString()); + Dataset counted = df.groupBy().count(); + + String plan = counted.queryExecution().executedPlan().toString(); + assertTrue( + plan.contains("VortexCountStarScan"), + "expected the physical plan to use the pushed-down count scan, got:\n" + plan); + assertEquals( + 10_000L, counted.collectAsList().get(0).getLong(0), "pushed-down count must match the written rows"); + assertEquals(10_000L, df.count(), "Dataset.count() must match the written rows"); + } + + @Test + @DisplayName("COUNT(*) with a data filter is not pushed down and stays correct") + public void testFilteredCountFallsBackToRegularScan() throws IOException { + Path outputPath = writeRows(1_000, "count_filtered", 2); + + Dataset filtered = spark.read() + .format("vortex") + .load(outputPath.toUri().toString()) + .filter("id < 100"); + Dataset counted = filtered.groupBy().count(); + + String plan = counted.queryExecution().executedPlan().toString(); + assertFalse( + plan.contains("VortexCountStarScan"), "a filtered count must not use footer metadata, got:\n" + plan); + assertEquals(100L, filtered.count(), "filtered count must reflect the predicate, not footer totals"); + } + + @Test + @DisplayName("Grouped counts are not pushed down and stay correct") + public void testGroupedCountFallsBackToRegularScan() throws IOException { + Path outputPath = writeRows(1_000, "count_grouped", 2); + + Dataset df = spark.read().format("vortex").load(outputPath.toUri().toString()); + Dataset grouped = + df.selectExpr("id % 2 as bucket").groupBy("bucket").count(); + + String plan = grouped.queryExecution().executedPlan().toString(); + assertFalse(plan.contains("VortexCountStarScan"), "a grouped count must scan data, got:\n" + plan); + assertEquals(2, grouped.collectAsList().size(), "grouping must produce one row per bucket"); + } + + @Test + @DisplayName("Builder accepts a single COUNT(*) aggregation and builds a VortexCountStarScan") + public void testBuilderAcceptsSingleCountStar() { + VortexScanBuilder builder = new VortexScanBuilder(Map.of()); + builder.addPath("/tmp/example.vortex"); + + assertTrue( + builder.pushAggregation(aggregation(new CountStar())), + "a lone COUNT(*) with no grouping must be pushed"); + assertInstanceOf( + VortexCountStarScan.class, builder.build(), "a pushed COUNT(*) must build the metadata count scan"); + } + + @Test + @DisplayName("Builder declines aggregations it cannot answer from footer metadata") + public void testBuilderDeclinesUnsupportedAggregations() { + assertFalse( + newBuilder().pushAggregation(aggregation(new Count(Expressions.column("id"), false))), + "COUNT(col) must not be pushed: footer counts do not reflect column nulls"); + assertFalse( + newBuilder().pushAggregation(aggregation(new Min(Expressions.column("id")))), "MIN must not be pushed"); + assertFalse( + newBuilder().pushAggregation(aggregation(new CountStar(), new CountStar())), + "multiple aggregate expressions must not be pushed"); + assertFalse( + newBuilder() + .pushAggregation(new Aggregation( + new AggregateFunc[] {new CountStar()}, new Expression[] {Expressions.column("id")})), + "grouped aggregations must not be pushed"); + } + + @Test + @DisplayName("Builder declining an aggregation leaves the regular scan path intact") + public void testDeclinedAggregationBuildsRegularScan() { + VortexScanBuilder builder = newBuilder(); + + assertFalse(builder.pushAggregation(aggregation(new Min(Expressions.column("id"))))); + assertInstanceOf(VortexScan.class, builder.build(), "a declined pushdown must build the regular scan"); + } + + private static VortexScanBuilder newBuilder() { + VortexScanBuilder builder = new VortexScanBuilder(Map.of()); + builder.addPath("/tmp/example.vortex"); + return builder; + } + + private static Aggregation aggregation(AggregateFunc... functions) { + return new Aggregation(functions, new Expression[0]); + } + + private Path writeRows(int numRows, String name, int partitions) throws IOException { + Path outputPath = tempDir.resolve(name); + Dataset df = spark.range(0, numRows) + .selectExpr("cast(id as int) as id", "concat('value_', cast(id as string)) as value"); + + df.repartition(partitions) + .write() + .format("vortex") + .option("path", outputPath.toUri().toString()) + .mode(SaveMode.Overwrite) + .save(); + return outputPath; + } +}