perf(calcite): push a filter on a narrow integer column through Calcite's cast

CHANGELOG.md

@@ -17,6 +17,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
+- A SQL `WHERE` comparing a `SMALLINT` or `TINYINT` column with a literal is pushed into the scan and its aggregates, instead of being left to Calcite because of the cast Calcite adds. ([#560](https://github.com/dfa1/vortex-java/pull/560))
 - SQL `MIN`/`MAX` of a date, time or timestamp column is answered from zone maps, and scans read those columns without building a `java.time` object per row. ([#559](https://github.com/dfa1/vortex-java/pull/559))
 - The Calcite adapter declares `TIME` and `TIMESTAMP` columns with millisecond precision, so `getString` shows the milliseconds the values carry. ([#559](https://github.com/dfa1/vortex-java/pull/559))
 - SQL `LIKE 'prefix%'` on a string column is about 5 times faster: it is pushed into the scan as a string range that prunes chunks, instead of being matched by Calcite row by row. ([#558](https://github.com/dfa1/vortex-java/pull/558))

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

@@ -13,6 +13,7 @@
 import org.apache.calcite.rex.RexNode;
 import org.apache.calcite.rex.RexUtil;
 import org.apache.calcite.sql.SqlKind;
+import org.apache.calcite.sql.type.SqlTypeName;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -268,10 +269,10 @@ private static Optional<RowFilter> binary(RexCall call, List<ColumnName> names,
         SqlKind kind = call.getKind();
         RexInputRef ref;
         RexLiteral lit;
-        if (ops.get(0) instanceof RexInputRef r && ops.get(1) instanceof RexLiteral l) {
+        if (columnRef(ops.get(0)) instanceof RexInputRef r && ops.get(1) instanceof RexLiteral l) {
             ref = r;
             lit = l;
-        } else if (ops.get(0) instanceof RexLiteral l && ops.get(1) instanceof RexInputRef r) {
+        } else if (ops.get(0) instanceof RexLiteral l && columnRef(ops.get(1)) instanceof RexInputRef r) {
             // `5 < x` is `x > 5`
             ref = r;
             lit = l;
@@ -296,6 +297,42 @@ private static Optional<RowFilter> binary(RexCall call, List<ColumnName> names,
         });
     }
 
+    /// The column a comparison operand reads: a bare reference, or one under a `CAST` that widens it
+    /// without loss, which Calcite adds when a narrow column meets a wider literal
+    /// (`order_serverId < 5` arrives as `CAST($5):INTEGER < 5` for a `SMALLINT` column). Comparing
+    /// the column itself is then the same comparison. `null` for anything else, including a `CAST`
+    /// that can change a value (`BIGINT` to `DOUBLE` rounds past 2^53).
+    private static RexInputRef columnRef(RexNode node) {
+        if (node instanceof RexInputRef ref) {
+            return ref;
+        }
+        if (node instanceof RexCall cast && cast.getKind() == SqlKind.CAST
+                && cast.getOperands().getFirst() instanceof RexInputRef ref
+                && widens(ref.getType().getSqlTypeName(), cast.getType().getSqlTypeName())) {
+            return ref;
+        }
+        return null;
+    }
+
+    private static boolean widens(SqlTypeName from, SqlTypeName to) {
+        int fromRank = integerRank(from);
+        int toRank = integerRank(to);
+        if (fromRank > 0 && toRank > 0) {
+            return toRank >= fromRank;
+        }
+        return (from == SqlTypeName.REAL || from == SqlTypeName.FLOAT) && to == SqlTypeName.DOUBLE;
+    }
+
+    private static int integerRank(SqlTypeName type) {
+        return switch (type) {
+            case TINYINT -> 1;
+            case SMALLINT -> 2;
+            case INTEGER -> 3;
+            case BIGINT -> 4;
+            default -> 0;
+        };
+    }
+
     /// Coerces a SQL literal to the Java type the column's zone-map statistics are stored as, so
     /// the comparison in [RowFilter] is type-compatible. Integer scalars are stored as `Long` and
     /// floats as `Double` in the stats, regardless of the column's physical width — a mismatched

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

@@ -100,6 +100,42 @@ void isNull_prunesChunksWithoutNulls() throws Exception {
         assertScan(isNull, List.of(5L), 1);
     }
 
+    /// Calcite casts a narrow column up to meet a wider literal (`CAST($0):INTEGER < 5` for a
+    /// `SMALLINT`), and the translation rejected anything but a bare column, so such a filter was
+    /// neither pushed nor enforced and every aggregate over it fell back to the scan.
+    @Test
+    void wideningCast_translatesAsTheColumn() {
+        // Given smallint < 5, as Calcite writes it
+        RelDataType smallint = REX.getTypeFactory().createSqlType(SqlTypeName.SMALLINT);
+        RelDataType integer = REX.getTypeFactory().createSqlType(SqlTypeName.INTEGER);
+        RexNode cast = REX.makeCast(integer, REX.makeInputRef(smallint, 0));
+        RexNode lessThan = REX.makeCall(SqlStdOperatorTable.LESS_THAN, cast,
+                REX.makeExactLiteral(BigDecimal.valueOf(5), integer));
+
+        // When
+        var result = RexFilterTranslator.translateStrict(List.of(lessThan),
+                List.of(ColumnName.of("small")), List.of(DType.I16));
+
+        // Then
+        assertThat(result).contains(io.github.dfa1.vortex.reader.RowFilter.lt(ColumnName.of("small"), 5L));
+    }
+
+    /// `BIGINT` to `DOUBLE` rounds past 2^53, so comparing the cast is not comparing the column.
+    @Test
+    void lossyCast_isLeftToCalcite() {
+        // Given bigint-as-double > 1.5
+        RelDataType dbl = REX.getTypeFactory().createSqlType(SqlTypeName.DOUBLE);
+        RexNode cast = REX.makeCast(dbl, REX.makeInputRef(BIGINT, 0));
+        RexNode greater = REX.makeCall(SqlStdOperatorTable.GREATER_THAN, cast,
+                REX.makeApproxLiteral(new BigDecimal("1.5"), dbl));
+
+        // When
+        var result = RexFilterTranslator.translateStrict(List.of(greater), List.of(ColumnName.of("id")), List.of(DType.I64));
+
+        // Then
+        assertThat(result).isEmpty();
+    }
+
     private static void assertScan(RexNode predicate, List<Long> expectedIds, long expectedChunks) {
         // When
         List<RexNode> filters = new ArrayList<>(List.of(predicate));

integration/src/test/java/io/github/dfa1/vortex/integration/CalciteInteropIntegrationTest.java

@@ -349,6 +349,56 @@ void jniWriter_temporalMinMax_answeredFromStatsEqualsTheScan(String aggregates,
         }
     }
 
+    @ParameterizedTest(name = "where {0}")
+    @CsvSource(delimiter = '|', value = {
+            "i16 < 5                 | i16 + 0 < 5",
+            "i16 >= -3 and i16 <> 7  | i16 + 0 >= -3 and i16 + 0 <> 7",
+            "2 > i16                 | 2 > i16 + 0"})
+    void jniWriter_aggregateWhereOnSmallint_isPushedDownAndEqualsTheScan(String where, String oracle, @TempDir Path tmp)
+            throws Exception {
+        // Given — Calcite compares a SMALLINT column through CAST(col AS INTEGER); the filter
+        // translation rejected the cast, so every aggregate filtered on such a column scanned
+        Path file = tmp.resolve("jni_smallint.vortex");
+        Schema schema = new Schema(List.of(
+                Field.notNullable("s", new ArrowType.Utf8()),
+                Field.notNullable("i16", new ArrowType.Int(16, true))));
+        String[] symbols = {"AAPL", "MSFT", "NVDA"};
+        try (VortexWriter writer = VortexWriter.builder(SESSION, file.toAbsolutePath().toUri().toString(), schema,
+                ALLOCATOR).build();
+             VectorSchemaRoot root = VectorSchemaRoot.create(schema, ALLOCATOR)) {
+            org.apache.arrow.vector.VarCharVector s = (org.apache.arrow.vector.VarCharVector) root.getVector("s");
+            org.apache.arrow.vector.SmallIntVector i16 = (org.apache.arrow.vector.SmallIntVector) root.getVector("i16");
+            s.allocateNew(BATCH_ROWS);
+            i16.allocateNew(BATCH_ROWS);
+            for (int i = 0; i < BATCH_ROWS; i++) {
+                s.setSafe(i, symbols[i % symbols.length].getBytes(java.nio.charset.StandardCharsets.UTF_8));
+                i16.set(i, (short) (i % 20 - 10));
+            }
+            root.setRowCount(BATCH_ROWS);
+            try (ArrowArray array = ArrowArray.allocateNew(ALLOCATOR);
+                 ArrowSchema arrowSchema = ArrowSchema.allocateNew(ALLOCATOR)) {
+                Data.exportVectorSchemaRoot(ALLOCATOR, root, null, array, arrowSchema);
+                writer.writeBatch(array.memoryAddress(), arrowSchema.memoryAddress());
+            }
+        }
+        String select = "select s, count(*), min(i16), max(i16) from vtx.t where ";
+        String sql = select + where + " group by s order by s";
+
+        try (Connection conn = VortexCalcite.connect("vtx", Map.of("t", file))) {
+            // When
+            String plan;
+            try (Statement st = conn.createStatement(); ResultSet rs = st.executeQuery("explain plan for " + sql)) {
+                rs.next();
+                plan = rs.getString(1);
+            }
+            List<String> result = sqlRows(conn, sql);
+
+            // Then
+            assertThat(plan).containsIgnoringCase("Values").doesNotContain("TableScan");
+            assertThat(result).isEqualTo(sqlRows(conn, select + oracle + " group by s order by s"));
+        }
+    }
+
     private static List<String> sqlRows(Connection conn, String sql) throws Exception {
         List<String> out = new java.util.ArrayList<>();
         try (Statement st = conn.createStatement(); ResultSet rs = st.executeQuery(sql)) {