Push down global COUNT(*) to footer row counts in vortex-spark (#9655)

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.
+ *
+ * <p>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<String> paths;
+    private final Map<String, String> 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<String> paths, Map<String, String> formatOptions) {
+        this.paths = paths;
+        this.formatOptions = formatOptions;
+    }
+
+    @Override
+    public InputPartition[] planInputPartitions() {
+        List<String> 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);
+    }
+}

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.
+ *
+ * <p>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 {}

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}.
+ *
+ * <p>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<ColumnarBatch> {
+    private final String path;
+    private final Map<String, String> 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<String, String> 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();
+    }
+}

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.
+ *
+ * <p>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<String, String> formatOptions;
+
+    /**
+     * Creates a new factory.
+     *
+     * @param formatOptions the format options for opening the files
+     */
+    VortexCountStarPartitionReaderFactory(Map<String, String> formatOptions) {
+        this.formatOptions = formatOptions;
+    }
+
+    @Override
+    public PartitionReader<InternalRow> createReader(InputPartition partition) {
+        throw new UnsupportedOperationException("row-based reads are not supported");
+    }
+
+    @Override
+    public PartitionReader<ColumnarBatch> 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;
+    }
+}

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.
+ *
+ * <p>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<String> paths;
+    private final Map<String, String> 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<String> paths, Map<String, String> 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;
+    }
+}

java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexScanBuilder.java

@@ -19,23 +19,28 @@
 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;
 import org.apache.spark.sql.types.StructType;
 
 /** 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<String> paths;
     private final List<Column> tableColumns;
     private final List<Column> readColumns;
     private final Map<String, String> formatOptions;
     private final Set<String> partitionColumnNames;
     private Predicate[] pushedPredicates = new Predicate[0];
+    private boolean pushedCountStar = false;
 
     /** Creates a new VortexScanBuilder with empty paths and columns. */
     public VortexScanBuilder(Map<String, String> 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.
+     *
+     * <p>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.
+     *
+     * <p>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.
      *
      * <p>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.
      *
+     * <p>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)));
     }