feat(calcite): read ORDER BY without LIMIT lazily on big tables

CHANGELOG.md

@@ -18,6 +18,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Custom extension types: register an `ExtensionDecoder` on `ReadRegistry` and an `ExtensionEncoder` on `WriteRegistry`, with a `validate` hook on both sides, and `Chunk.as` decodes any registered type. ([#530](https://github.com/dfa1/vortex-java/issues/530))
 
 ### Changed
+- SQL `ORDER BY` with no `LIMIT` over a large table returns its first rows without sorting the whole table first. ([#570](https://github.com/dfa1/vortex-java/pull/570))
 - **Breaking:** `VortexWriter.create` refuses a nullable root struct: Rust reads a nullable struct only with a validity child the writer never wrote, so such files were unreadable outside vortex-java. ([#493](https://github.com/dfa1/vortex-java/issues/493))
 - **Breaking:** a file that does not embed its dtype is refused when it opens, as Rust refuses it unless the caller supplies the dtype, instead of failing the scan with `NullPointerException`. ([#493](https://github.com/dfa1/vortex-java/issues/493))
 - **Breaking:** every FlatBuffer in a file is verified before it is read, as Rust verifies it, so corrupt file metadata fails with `VortexException`, and like Rust the reader now refuses tables nested more than 64 deep. ([#493](https://github.com/dfa1/vortex-java/issues/493))

calcite/src/main/java/io/github/dfa1/vortex/calcite/VortexOrderedTable.java

@@ -0,0 +1,137 @@
+package io.github.dfa1.vortex.calcite;
+
+import io.github.dfa1.vortex.reader.RowFilter;
+
+import org.apache.calcite.DataContext;
+import org.apache.calcite.linq4j.AbstractEnumerable;
+import org.apache.calcite.linq4j.Enumerable;
+import org.apache.calcite.linq4j.Enumerator;
+import org.apache.calcite.rel.RelCollation;
+import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.rel.type.RelDataTypeFactory;
+import org.apache.calcite.schema.ScannableTable;
+import org.apache.calcite.schema.Statistic;
+import org.apache.calcite.schema.Statistics;
+import org.apache.calcite.schema.impl.AbstractTable;
+
+import java.util.List;
+
+/// A [VortexTable] already sorted: what `ORDER BY` with no `LIMIT` over a big table is rewritten to
+/// scan, so the rows come back lazily instead of through Calcite's sort, which buffers every row of
+/// the table before it returns the first (a 100M-row table ran an IDE's JDBC process out of memory).
+///
+/// The first rows come from a bounded top-N pass, [VortexTable#topN]. When the client reads past
+/// them, the pass runs again for eight times as many rows and the rows already returned are skipped.
+/// Ties rank by scan position, so a longer pass always starts with the rows of the shorter one.
+///
+/// Each pass reads the whole table: the heap bounds memory, not I/O. A client that stops after a grid
+/// page costs one pass and holds only that page, which is the case this exists for. One that reads on
+/// pays another full pass each time it outgrows the last (a 100M-row table read to the end is seven), and holds what Calcite's own sort would at the end. A file replaced while a result is being
+/// read is not detected: each pass reopens it by path.
+final class VortexOrderedTable extends AbstractTable implements ScannableTable {
+
+    /// Rows the first pass keeps: a few IDE grid pages.
+    private static final int FIRST_BATCH = 1024;
+
+    /// How much longer each later pass is. Every pass is a full scan, so growing fast bounds how many a
+    /// client that reads far pays for.
+    private static final int GROWTH = 8;
+
+    private final VortexTable table;
+    private final int[] cols;
+    private final RowFilter filter;
+    private final List<VortexTable.SortKey> keys;
+    private final RelDataType rowType;
+    private final double estimatedRows;
+    private final RelCollation collation;
+
+    VortexOrderedTable(VortexTable table, int[] cols, RowFilter filter, List<VortexTable.SortKey> keys,
+                       RelDataType rowType, double estimatedRows, RelCollation collation) {
+        this.table = table;
+        this.cols = cols;
+        this.filter = filter;
+        this.keys = keys;
+        this.rowType = rowType;
+        this.estimatedRows = estimatedRows;
+        this.collation = collation;
+    }
+
+    @Override
+    public RelDataType getRowType(RelDataTypeFactory typeFactory) {
+        return rowType;
+    }
+
+    @Override
+    public Statistic getStatistic() {
+        // the planner's estimate for the rows the sort would have received (a filter shrinks it), so this
+        // scan is costed against the plain sort over the same filtered scan on equal terms; the
+        // enumerator below never uses it as a bound
+        // declaring the order is what lets the planner drop the sort the query asked for: without it
+        // it adds one back on top and the cheaper plain scan wins
+        return Statistics.of(estimatedRows, List.of(), List.of(collation));
+    }
+
+    @Override
+    public Enumerable<Object[]> scan(DataContext root) {
+        return new AbstractEnumerable<>() {
+            @Override
+            public Enumerator<Object[]> enumerator() {
+                return new Batches();
+            }
+        };
+    }
+
+    /// Returns the sorted rows, running a longer top-N pass whenever the current one is used up.
+    private final class Batches implements Enumerator<Object[]> {
+
+        private List<Object[]> batch = List.of();
+        private int requested;
+        private long returned;
+        private int index = -1;
+        private boolean exhausted;
+
+        @Override
+        public boolean moveNext() {
+            index++;
+            if (index < batch.size()) {
+                returned = index + 1L;
+                return true;
+            }
+            if (exhausted) {
+                return false;
+            }
+            // the old batch is dead weight while the next pass runs
+            batch = List.of();
+            // the first pass keeps FIRST_BATCH rows, every later one GROWTH times the last; a pass
+            // cannot keep more rows than the table has
+            long tableRows = (long) table.getStatistic().getRowCount().doubleValue();
+            long want = requested == 0 ? FIRST_BATCH : (long) GROWTH * requested;
+            requested = (int) Math.min(want, Math.min(Integer.MAX_VALUE - 8L, Math.max(1, tableRows)));
+            batch = table.topN(cols, filter, keys, requested);
+            exhausted = batch.size() < requested || requested >= tableRows;
+            // skip the rows an earlier, shorter pass already returned; counted when advanced over, not
+            // when read, since a client may step through rows without reading a column
+            index = (int) returned;
+            boolean more = index < batch.size();
+            if (more) {
+                returned = index + 1L;
+            }
+            return more;
+        }
+
+        @Override
+        public Object[] current() {
+            return batch.get(index);
+        }
+
+        @Override
+        public void reset() {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public void close() {
+            batch = List.of();
+        }
+    }
+}

calcite/src/main/java/io/github/dfa1/vortex/calcite/VortexSortPushDownRule.java

@@ -7,12 +7,15 @@
 import org.apache.calcite.plan.RelOptRule;
 import org.apache.calcite.plan.RelOptRuleCall;
 import org.apache.calcite.plan.RelOptRuleOperand;
+import org.apache.calcite.plan.RelOptTable;
+import org.apache.calcite.prepare.RelOptTableImpl;
 import org.apache.calcite.rel.RelFieldCollation;
 import org.apache.calcite.rel.RelNode;
 import org.apache.calcite.rel.core.Filter;
 import org.apache.calcite.rel.core.Project;
 import org.apache.calcite.rel.core.Sort;
 import org.apache.calcite.rel.core.TableScan;
+import org.apache.calcite.rel.logical.LogicalTableScan;
 import org.apache.calcite.rel.logical.LogicalValues;
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.rex.RexBuilder;
@@ -49,6 +52,10 @@ public final class VortexSortPushDownRule extends RelOptRule {
     /// Most rows a rewrite emits; each becomes a literal row of the plan.
     private static final int MAX_ROWS = 1000;
 
+    /// Smallest table, in rows, whose `ORDER BY` with no `LIMIT` is rewritten to a lazy scan
+    /// ([VortexOrderedTable]); a smaller one fits in memory and Calcite sorts it faster.
+    static final int LAZY_SORT_MIN_ROWS = 100_000;
+
     /// Most literal cells (rows times columns) a push-down may put in one `VALUES`. Calcite compiles
     /// the whole `VALUES` into a single Java method, which the JVM caps at 64 KB of bytecode: measured
     /// on a mixed-type table, 4,500 cells compile and 6,000 do not ("Code grows beyond 64 KB"). A bigger
@@ -73,14 +80,24 @@ private VortexSortPushDownRule(RelOptRuleOperand operand, String description) {
     @Override
     public void onMatch(RelOptRuleCall call) {
         Sort sort = call.rel(0);
-        if (sort.fetch == null || sort.getCollation().getFieldCollations().isEmpty()) {
-            return; // a full ORDER BY must emit every row
-        }
-        int fetch = intLiteral(sort.fetch);
-        int offset = sort.offset == null ? 0 : intLiteral(sort.offset);
-        if (fetch < 0 || offset < 0 || fetch + (long) offset > MAX_ROWS) {
+        if (sort.getCollation().getFieldCollations().isEmpty()) {
             return;
         }
+        // no LIMIT: every row must come back sorted, which only a big table's lazy scan can do cheaply
+        boolean lazy = sort.fetch == null;
+        int fetch = 0;
+        int offset = 0;
+        if (lazy) {
+            if (sort.offset != null) {
+                return;
+            }
+        } else {
+            fetch = intLiteral(sort.fetch);
+            offset = sort.offset == null ? 0 : intLiteral(sort.offset);
+            if (fetch < 0 || offset < 0 || fetch + (long) offset > MAX_ROWS) {
+                return;
+            }
+        }
         Project project = null;
         Filter filter = null;
         RelNode current = call.rel(1);
@@ -98,6 +115,10 @@ public void onMatch(RelOptRuleCall call) {
         if (table == null) {
             return;
         }
+        Double rowCount = table.getStatistic().getRowCount();
+        if (lazy && (rowCount == null || rowCount <= LAZY_SORT_MIN_ROWS)) {
+            return; // small enough for Calcite to sort in memory
+        }
         List<RexNode> predicates = new ArrayList<>();
         if (filter != null) {
             predicates.add(filter.getCondition());
@@ -115,7 +136,7 @@ public void onMatch(RelOptRuleCall call) {
         }
         int[] cols = columns(scan, project);
         // only the rows after the offset become literals; the offset rows stay in the heap (MAX_ROWS)
-        if (cols == null || (long) fetch * cols.length > MAX_VALUES_CELLS) {
+        if (cols == null || (!lazy && (long) fetch * cols.length > MAX_VALUES_CELLS)) {
             return;
         }
         List<VortexTable.SortKey> keys = new ArrayList<>();
@@ -126,6 +147,11 @@ public void onMatch(RelOptRuleCall call) {
             keys.add(new VortexTable.SortKey(collation.getFieldIndex(), descending,
                     nulls == RelFieldCollation.NullDirection.FIRST));
         }
+        if (lazy) {
+            call.transformTo(orderedScan(sort, scan, new VortexOrderedTable(table, cols, rowFilter, keys,
+                    sort.getRowType(), call.getMetadataQuery().getRowCount(sort.getInput()), sort.getCollation())));
+            return;
+        }
         int count = offset + fetch;
         RowFilter finalFilter = rowFilter;
         List<Object> foldKey = List.of("topN", table, java.util.Arrays.toString(cols),
@@ -150,6 +176,15 @@ public void onMatch(RelOptRuleCall call) {
         call.transformTo(LogicalValues.create(sort.getCluster(), sort.getRowType(), ImmutableList.copyOf(tuples)));
     }
 
+    /// A scan of `ordered` that Calcite plans like any table, in place of the sort and everything it
+    /// sorted.
+    private static RelNode orderedScan(Sort sort, TableScan scan, VortexOrderedTable ordered) {
+        RelOptTable base = scan.getTable();
+        RelOptTable derived = RelOptTableImpl.create(base.getRelOptSchema(), sort.getRowType(), ordered,
+                ImmutableList.<String>builder().addAll(base.getQualifiedName()).add("ordered").build());
+        return LogicalTableScan.create(sort.getCluster(), derived, scan.getHints());
+    }
+
     /// The table columns the sort's input emits, in order, or `null` when a projected expression is
     /// not a plain column.
     private static int[] columns(TableScan scan, Project project) {

calcite/src/main/java/io/github/dfa1/vortex/calcite/VortexTable.java

@@ -645,7 +645,7 @@ public List<Object[]> topN(int[] cols, RowFilter filter, List<SortKey> keys, int
                 .thenComparingLong(Ranked::position);
         // the worst kept row at the head, so a better row evicts it
         java.util.PriorityQueue<Ranked> heap = new java.util.PriorityQueue<>(count + 1, ranking.reversed());
-        VortexEnumerator rows = new VortexEnumerator(file, new AtomicLong(), options, outNames, outTypes, filter);
+        VortexEnumerator rows = new VortexEnumerator(file, chunksScannedLastQuery, options, outNames, outTypes, filter);
         // only the key columns are read per row; a row is built once it beats the worst one kept. A
         // later row never beats an equal key, so the first rows of a tie are the ones kept
         Object[] probe = new Object[cols.length];