fix(calcite): keep push-down VALUES under the JVM's 64 KB method limit

CHANGELOG.md

@@ -44,6 +44,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - **Breaking:** `ExtensionId` is a sealed `WellKnown`/`Custom` type like `EncodingId`, so third-party extension ids can be named, and `ExtensionId.parse` returns the id instead of an `Optional`. ([#530](https://github.com/dfa1/vortex-java/issues/530))
 
 ### Fixed
+- SQL `ORDER BY ... LIMIT` and `GROUP BY` with a few thousand result cells no longer fail with "Code grows beyond 64 KB"; past the limit Calcite scans instead. ([#569](https://github.com/dfa1/vortex-java/pull/569))
 - A file whose array FlatBuffer holds a corrupt offset fails the read with `VortexException` instead of `IndexOutOfBoundsException`. ([#493](https://github.com/dfa1/vortex-java/issues/493))
 - A column whose dtype does not fit its bit-packed, delta or RLE encoding fails with `VortexException` instead of `ClassCastException`, and an empty RLE boolean chunk reads. ([#493](https://github.com/dfa1/vortex-java/issues/493))
 - A struct dtype missing its field names or types, or a struct layout whose children do not match its fields, is rejected when the file opens, as Rust does, instead of failing the scan with `ArrayIndexOutOfBoundsException`. ([#493](https://github.com/dfa1/vortex-java/issues/493))

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

@@ -190,11 +190,12 @@ public void onMatch(RelOptRuleCall call) {
     }
 
     /// Rewrites a `GROUP BY` on one string, integer or boolean column into a [LogicalValues] with one
-    /// row per group, folded by [VortexTable#groupedFold(ColumnName, RowFilter, java.util.Set)]
+    /// row per group, folded by [VortexTable#groupedFold(ColumnName, RowFilter, java.util.Set, int)]
     /// instead of materializing every row for Calcite to hash. Leaves the plan alone when the group
     /// column is not a string, an aggregate is not `COUNT`/`SUM`/`MIN`/`MAX` over a plain column, or
     /// the fold declines (a chunk not dictionary-encoded, an aggregate column it cannot total, too
-    /// many groups).
+    /// many groups) or the groups would not fit one compiled `VALUES`
+    /// (`VortexSortPushDownRule.MAX_VALUES_CELLS`).
     private static void rewriteGrouped(RelOptRuleCall call, Aggregate aggregate, VortexTable table,
                                        List<ColumnName> scanColumns, RelDataType scanRowType, Project project,
                                        RexBuilder rexBuilder, List<RelDataType> outTypes, RowFilter filter) {
@@ -227,10 +228,14 @@ private static void rewriteGrouped(RelOptRuleCall call, Aggregate aggregate, Vor
             callColumns.add(column);
             aggColumns.add(column);
         }
+        // every group is one literal row of a single compiled VALUES, so the fold gives up as soon as the
+        // groups would not fit, instead of scanning the whole file to find that out
+        int maxGroups = VortexSortPushDownRule.MAX_VALUES_CELLS / (1 + calls.size());
         // ConcurrentHashMap keys cannot be null, so a missing filter is keyed as an empty list
-        List<Object> foldKey = List.of("grouped", table, group, filter == null ? List.of() : filter, aggColumns);
+        List<Object> foldKey = List.of("grouped", table, group, filter == null ? List.of() : filter, aggColumns,
+                maxGroups);
         Optional<Map<Object, VortexTable.GroupFold>> folds =
-                memo(call.getPlanner(), foldKey, () -> table.groupedFold(group, filter, aggColumns));
+                memo(call.getPlanner(), foldKey, () -> table.groupedFold(group, filter, aggColumns, maxGroups));
         if (folds.isEmpty()) {
             return;
         }

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

@@ -49,6 +49,12 @@ 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;
 
+    /// 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
+    /// result is left to Calcite's own scan and sort, which are slower but correct.
+    static final int MAX_VALUES_CELLS = 4_000;
+
     /// Every rule variant, for registering with a planner in one call.
     public static final List<RelOptRule> RULES = List.of(
             new VortexSortPushDownRule(operand(Sort.class, operand(Project.class,
@@ -108,7 +114,8 @@ public void onMatch(RelOptRuleCall call) {
             rowFilter = translated.get();
         }
         int[] cols = columns(scan, project);
-        if (cols == null) {
+        // 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) {
             return;
         }
         List<VortexTable.SortKey> keys = new ArrayList<>();

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

@@ -362,11 +362,7 @@ public Optional<FilteredFold> filteredFold(RowFilter filter, ColumnName aggColum
         }
     }
 
-    /// Most groups [#groupedFold(ColumnName, RowFilter, Set)] answers. Each group becomes one literal
-    /// row of the plan, so a high-cardinality column is left to the scan.
-    private static final int MAX_GROUPS = 10_000;
-
-    /// One group's fold, returned by [#groupedFold(ColumnName, RowFilter, Set)].
+    /// One group's fold, returned by [#groupedFold(ColumnName, RowFilter, Set, int)].
     ///
     /// @param rows    the rows of the group the filter selects
     /// @param columns per aggregate column, its fold over those rows
@@ -380,14 +376,16 @@ public record GroupFold(long rows, Map<ColumnName, FilteredFold> columns) {
     ///
     /// The result is empty, leaving the query to the scan, when the kernel declines a chunk (a string
     /// column not dictionary-encoded there, too many groups in it, an aggregate column with no long
-    /// or double fold) or there are more than [#MAX_GROUPS] groups in all.
+    /// or double fold) or there are more than `maxGroups` groups in all.
     ///
     /// @param group      the column to group by
     /// @param filter     the pushed predicate, or `null` for every row
     /// @param aggColumns the columns the aggregates reduce; empty when only rows are counted
+    /// @param maxGroups  most groups to answer; the scan stops as soon as it sees more
     /// @return the fold per group value (`String`, `Boolean` or `Long`; `null` is the null group), or
     ///         empty
-    public Optional<Map<Object, GroupFold>> groupedFold(ColumnName group, RowFilter filter, Set<ColumnName> aggColumns) {
+    public Optional<Map<Object, GroupFold>> groupedFold(ColumnName group, RowFilter filter, Set<ColumnName> aggColumns,
+                                                        int maxGroups) {
         java.util.LinkedHashSet<ColumnName> columns = new java.util.LinkedHashSet<>();
         columns.add(group);
         if (filter != null) {
@@ -431,7 +429,7 @@ public Optional<Map<Object, GroupFold>> groupedFold(ColumnName group, RowFilter
                             }
                         }
                     }
-                    if (rows.size() > MAX_GROUPS) {
+                    if (rows.size() > maxGroups) {
                         return Optional.empty();
                     }
                 }

calcite/src/test/java/io/github/dfa1/vortex/calcite/ValuesCodegenLimitTest.java

@@ -0,0 +1,192 @@
+package io.github.dfa1.vortex.calcite;
+
+import io.github.dfa1.vortex.core.model.ColumnName;
+import io.github.dfa1.vortex.core.model.DType;
+import io.github.dfa1.vortex.core.model.PType;
+import io.github.dfa1.vortex.writer.VortexWriter;
+import io.github.dfa1.vortex.writer.WriteOptions;
+
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+
+import java.nio.channels.FileChannel;
+import java.nio.file.Path;
+import java.nio.file.StandardOpenOption;
+import java.sql.Connection;
+import java.sql.ResultSet;
+import java.sql.Statement;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.stream.LongStream;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/// The sort and grouped push-downs answer with a literal `VALUES`, which Calcite compiles into one
+/// Java method: past about 5,000 cells the JVM refuses it ("Code grows beyond 64 KB") and the query
+/// fails outright. Past [VortexSortPushDownRule#MAX_VALUES_CELLS] the rules must leave the query to
+/// Calcite's own scan, which returns the same rows; below it they must still push down.
+class ValuesCodegenLimitTest {
+
+    private static final int GROUPS = 3_000;
+
+    @TempDir
+    static Path tmp;
+    private static Path file;
+
+    /// 3,000 groups of 3 rows and six columns, so that `count(*)` per group is 6,000 cells and
+    /// `ORDER BY ... LIMIT 800` is 4,800: both over the cap, yet each failed in the field at that size.
+    @BeforeAll
+    static void writeFile() throws Exception {
+        file = tmp.resolve("wide.vortex");
+        DType.Struct schema = DType.structBuilder()
+                .field(ColumnName.of("k"), new DType.Utf8(false))
+                .field(ColumnName.of("v"), DType.I64)
+                .field(ColumnName.of("a"), new DType.Primitive(PType.F64, false))
+                .field(ColumnName.of("b"), new DType.Primitive(PType.F64, false))
+                .field(ColumnName.of("c"), new DType.Primitive(PType.F64, false))
+                .field(ColumnName.of("d"), new DType.Primitive(PType.F64, false))
+                .build();
+        int rows = GROUPS * 3;
+        String[] k = new String[rows];
+        long[] v = new long[rows];
+        double[] a = new double[rows];
+        for (int i = 0; i < rows; i++) {
+            k[i] = "key-" + (i % GROUPS);
+            v[i] = i;
+            a[i] = i * 0.5;
+        }
+        try (var channel = FileChannel.open(file, StandardOpenOption.CREATE, StandardOpenOption.WRITE);
+             var writer = VortexWriter.create(channel, schema, WriteOptions.defaults())) {
+            writer.writeChunk(Map.of(ColumnName.of("k"), k, ColumnName.of("v"), v, ColumnName.of("a"), a,
+                    ColumnName.of("b"), a, ColumnName.of("c"), a, ColumnName.of("d"), a));
+        }
+    }
+
+    @Test
+    void groupedOverTheCap_fallsBackToCalciteAndReturnsEveryGroup() throws Exception {
+        // Given 3,000 groups x (key + count + sum) = 9,000 cells, which used to fail to compile
+        String sql = "select k, count(*), sum(v) from vtx.t group by k";
+        try (Connection conn = VortexCalcite.connect("vtx", Map.of("t", file))) {
+
+            // When
+            String plan = plan(conn, sql);
+            long groups = 0;
+            long total = 0;
+            try (Statement st = conn.createStatement(); ResultSet rs = st.executeQuery(sql)) {
+                while (rs.next()) {
+                    groups++;
+                    total += rs.getLong(2);
+                }
+            }
+
+            // Then Calcite scanned instead, and the answer is complete
+            assertThat(plan).contains("TableScan").doesNotContain("Values");
+            assertThat(groups).isEqualTo(GROUPS);
+            assertThat(total).isEqualTo(GROUPS * 3L);
+        }
+    }
+
+    @Test
+    void groupedUnderTheCap_isStillPushedDown() throws Exception {
+        // Given a filter leaving 300 groups x 3 cells = 900 cells, well under the cap
+        String sql = "select k, count(*), sum(v) from vtx.t where v < 900 group by k";
+        try (Connection conn = VortexCalcite.connect("vtx", Map.of("t", file))) {
+
+            // When
+            String plan = plan(conn, sql);
+
+            // Then
+            assertThat(plan).contains("Values").doesNotContain("TableScan");
+        }
+    }
+
+    @ParameterizedTest
+    @ValueSource(ints = {700, 800, 1000})
+    void sortOverTheCap_fallsBackToCalciteAndReturnsTheTopRows(int limit) throws Exception {
+        // Given 6 columns x 700+ rows > 4,000 cells
+        String sql = "select * from vtx.t order by v limit " + limit;
+        try (Connection conn = VortexCalcite.connect("vtx", Map.of("t", file))) {
+
+            // When
+            String plan = plan(conn, sql);
+            long rows = 0;
+            long last = -1;
+            try (Statement st = conn.createStatement(); ResultSet rs = st.executeQuery(sql)) {
+                while (rs.next()) {
+                    rows++;
+                    last = rs.getLong("v");
+                }
+            }
+
+            // Then
+            assertThat(plan).doesNotContain("Values");
+            assertThat(rows).isEqualTo(limit);
+            assertThat(last).isEqualTo(limit - 1L);
+        }
+    }
+
+    @Test
+    void sortUnderTheCap_isStillPushedDown() throws Exception {
+        // Given 6 columns x 600 rows = 3,600 cells
+        String sql = "select * from vtx.t order by v limit 600";
+        try (Connection conn = VortexCalcite.connect("vtx", Map.of("t", file))) {
+
+            // When
+            String plan = plan(conn, sql);
+
+            // Then
+            assertThat(plan).contains("Values");
+        }
+    }
+
+    @Test
+    void pagedSortWithALargeOffset_isStillPushedDownWhenItsLiteralsFit() throws Exception {
+        // Given 100 rows x 6 columns = 600 literal cells; only the rows after the offset become
+        // literals, so the 900 skipped rows (6,000 cells counted naively) must not block the push-down
+        String sql = "select * from vtx.t order by v limit 100 offset 900";
+        try (Connection conn = VortexCalcite.connect("vtx", Map.of("t", file))) {
+
+            // When
+            String plan = plan(conn, sql);
+            List<Long> result = new ArrayList<>();
+            try (Statement st = conn.createStatement(); ResultSet rs = st.executeQuery(sql)) {
+                while (rs.next()) {
+                    result.add(rs.getLong("v"));
+                }
+            }
+
+            // Then
+            assertThat(plan).contains("Values");
+            assertThat(result).isEqualTo(LongStream.range(900, 1000).boxed().toList());
+        }
+    }
+
+    @Test
+    void groupedFold_givesUpAsSoonAsTheGroupsExceedTheLimit() {
+        // Given 3,000 groups
+        VortexTable sut = new VortexTable(file);
+        ColumnName key = ColumnName.of("k");
+        Set<ColumnName> aggregates = Set.of(ColumnName.of("v"));
+
+        // When
+        var tooMany = sut.groupedFold(key, null, aggregates, 2_000);
+        var enough = sut.groupedFold(key, null, aggregates, 3_000);
+
+        // Then
+        assertThat(tooMany).isEmpty();
+        assertThat(enough).isPresent();
+        assertThat(enough.get()).hasSize(GROUPS);
+    }
+
+    private static String plan(Connection conn, String sql) throws Exception {
+        try (Statement st = conn.createStatement(); ResultSet rs = st.executeQuery("explain plan for " + sql)) {
+            rs.next();
+            return rs.getString(1);
+        }
+    }
+}