Honour Calcite TIMESTAMP precision in QTF stitcher output schema

ArrowCalciteTypes.toArrow() hardcoded TimeUnit.MILLISECOND for every Calcite
TIMESTAMP, so the late-materialization Stitcher pre-allocated a
TimestampMilliVector for the output VSR. Shards correctly emitted
Timestamp(NANOSECOND) for date_nanos fields, and Arrow's
BaseFixedWidthVector.copyFromSafe requires identical fixed-width subclasses,
so it tripped on every shard's response with a message-less
IllegalArgumentException — surfacing to clients as HTTP 400 "Invalid Query".

Plan shape required to fire: multi-shard parquet-backed index with date_nanos +
Timestamp in projection + sort + head N (or other top-K) + at least one
fetch-only field above the anchor. Drop any one and the rewriter declines QTF
or the bug doesn't fire.

Fix: branch on Calcite precision. precision == 9 -> NANOSECOND, else MILLISECOND.
BackendPlanAdapter preserves precision through CAST rewrites, so the value is
reliable (verified via DEBUG plan dump showing TIMESTAMP(9) survives).

- Unit: ArrowCalciteTypesTests adds 3 cases covering precision 3, 9, default.
  The factory uses an extended RelDataTypeSystemImpl that lifts MAX_DATETIME
  precision to 9 — default Calcite caps at 3 and would silently clamp the test
  input.
- IT: LateMaterializationDateNanosIT spins a 2-shard parquet+lucene index with
  date_nanos and runs the q01-style plan (match + sort - ts + head + fetch-only
  field). Pre-fix: HTTP 400 with the exact bug signature. Post-fix: 4 rows in
  DESC order with sub-millisecond precision preserved.

Signed-off-by: Vinay Krishna Pudyodu <vinkrish.neo@gmail.com>

Author: vinaykpud

sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/ArrowCalciteTypes.java

@@ -22,12 +22,6 @@
  * schema. The {@code AggregateFunction.ArrowToCalciteTypeMapper} (in the SPI module) handles
  * the inverse direction for {@code IntermediateField} resolution; this class is kept as the
  * single authority for the Calcite→Arrow direction needed outside that resolver.
- *
- * <p>FIXME [FixBeforeMainMerge] coverage gaps: TIMESTAMP currently hardcodes MILLISECOND
- * (Calcite precision is ignored — see toArrow note), so date_nanos is not yet distinguished,
- * and several Calcite/Arrow types are still unmapped (TIMESTAMP_WITH_LOCAL_TIME_ZONE, DATE,
- * TIME, DECIMAL, Arrow Date/Time/Decimal/...). Audit and broaden before merge so the QTF path
- * tolerates non-keyword/non-date columns end-to-end.
  */
 public final class ArrowCalciteTypes {
 
@@ -53,11 +47,8 @@ public static ArrowType toArrow(RelDataType t) {
             case VARBINARY, BINARY -> ArrowType.Binary.INSTANCE;
             case BOOLEAN -> ArrowType.Bool.INSTANCE;
             // TODO: TIMESTAMP_WITH_LOCAL_TIME_ZONE, DATE, TIME, DECIMAL still missing.
-            // TODO: hardcoded MILLISECOND to match what DateParquetField emits on the data node;
-            // Calcite's reported precision doesn't track the wire-level Arrow precision today, so
-            // honouring t.getPrecision() here would break Stitcher copyFromSafe. Revisit when
-            // Calcite types carry the data-node-side Arrow precision faithfully.
-            case TIMESTAMP -> new ArrowType.Timestamp(TimeUnit.MILLISECOND, null);
+            // precision 9 ⇒ date_nanos; else date — must match the wire unit shards emit (Stitcher copyFromSafe).
+            case TIMESTAMP -> new ArrowType.Timestamp(t.getPrecision() == 9 ? TimeUnit.NANOSECOND : TimeUnit.MILLISECOND, null);
             default -> throw new IllegalArgumentException("Unsupported Calcite type: " + t.getSqlTypeName());
         };
     }

sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/ArrowCalciteTypesTests.java

@@ -8,9 +8,11 @@
 
 package org.opensearch.analytics.planner;
 
+import org.apache.arrow.vector.types.TimeUnit;
 import org.apache.arrow.vector.types.pojo.ArrowType;
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.rel.type.RelDataTypeSystem;
+import org.apache.calcite.rel.type.RelDataTypeSystemImpl;
 import org.apache.calcite.sql.type.SqlTypeFactoryImpl;
 import org.apache.calcite.sql.type.SqlTypeName;
 import org.opensearch.test.OpenSearchTestCase;
@@ -23,10 +25,25 @@ public class ArrowCalciteTypesTests extends OpenSearchTestCase {
 
     private static final SqlTypeFactoryImpl TYPE_FACTORY = new SqlTypeFactoryImpl(RelDataTypeSystem.DEFAULT);
 
+    /** Lifts TIMESTAMP max-precision to 9; default Calcite caps at 3 and would clamp date_nanos away. */
+    private static final SqlTypeFactoryImpl NANOS_TYPE_FACTORY = new SqlTypeFactoryImpl(new RelDataTypeSystemImpl() {
+        @Override
+        public int getMaxPrecision(SqlTypeName typeName) {
+            if (typeName == SqlTypeName.TIMESTAMP || typeName == SqlTypeName.TIMESTAMP_WITH_LOCAL_TIME_ZONE) {
+                return 9;
+            }
+            return super.getMaxPrecision(typeName);
+        }
+    });
+
     private static RelDataType type(SqlTypeName name) {
         return TYPE_FACTORY.createSqlType(name);
     }
 
+    private static RelDataType timestamp(int precision) {
+        return NANOS_TYPE_FACTORY.createSqlType(SqlTypeName.TIMESTAMP, precision);
+    }
+
     /**
      * SMALLINT (OpenSearch {@code short}) must map to the wire Arrow type the data node
      * emits — {@code Int(16, true)} per {@code ShortParquetField} — so the Stitcher's
@@ -45,4 +62,19 @@ public void testIntegerAndBigintUnchanged() {
         assertEquals(new ArrowType.Int(32, true), ArrowCalciteTypes.toArrow(type(SqlTypeName.INTEGER)));
         assertEquals(new ArrowType.Int(64, true), ArrowCalciteTypes.toArrow(type(SqlTypeName.BIGINT)));
     }
+
+    /** date ⇒ TIMESTAMP(3) ⇒ MILLISECOND. */
+    public void testTimestampPrecision3MapsToMillisecond() {
+        assertEquals(new ArrowType.Timestamp(TimeUnit.MILLISECOND, null), ArrowCalciteTypes.toArrow(timestamp(3)));
+    }
+
+    /** date_nanos ⇒ TIMESTAMP(9) ⇒ NANOSECOND — regression: previously hardcoded MILLISECOND, tripped Stitcher copyFromSafe. */
+    public void testTimestampPrecision9MapsToNanosecond() {
+        assertEquals(new ArrowType.Timestamp(TimeUnit.NANOSECOND, null), ArrowCalciteTypes.toArrow(timestamp(9)));
+    }
+
+    /** Default-precision TIMESTAMP (precision 0) keeps the legacy MILLISECOND mapping. */
+    public void testTimestampDefaultPrecisionMapsToMillisecond() {
+        assertEquals(new ArrowType.Timestamp(TimeUnit.MILLISECOND, null), ArrowCalciteTypes.toArrow(type(SqlTypeName.TIMESTAMP)));
+    }
 }

sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/LateMaterializationDateNanosIT.java

@@ -0,0 +1,107 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.analytics.qa;
+
+import org.opensearch.client.Request;
+import org.opensearch.client.Response;
+
+import java.util.List;
+import java.util.Map;
+
+/** End-to-end test for late-materialization queries against a {@code date_nanos} timestamp on a multi-shard parquet index. */
+public class LateMaterializationDateNanosIT extends AnalyticsRestTestCase {
+
+    private static final String INDEX = "late_mat_date_nanos_e2e";
+    private static final int NUM_SHARDS = 2;
+
+    public void testMultiShardLateMatSortFetchOnDateNanos() throws Exception {
+        createParquetBackedIndex();
+        indexDocs();
+
+        // `severity` is fetch-only (projected, not in filter/sort) — required for the LM rewriter
+        // to fire; otherwise above ⊆ below and it skips, masking the bug.
+        Map<String, Object> result = executePpl(
+            "source = " + INDEX
+                + " | where match(body, 'failed') and service = 'checkout'"
+                + " | sort - ts"
+                + " | fields ts, severity, body"
+                + " | head 4"
+        );
+
+        List<String> columns = extractColumnNames(result);
+        assertTrue("schema must contain 'ts', got " + columns, columns.contains("ts"));
+        assertTrue("schema must contain 'severity', got " + columns, columns.contains("severity"));
+        assertTrue("schema must contain 'body', got " + columns, columns.contains("body"));
+
+        @SuppressWarnings("unchecked")
+        List<List<Object>> rows = (List<List<Object>>) result.get("datarows");
+        assertNotNull("Stitcher copyFromSafe trip returns no payload", rows);
+        assertEquals("4 matching docs across 2 shards", 4, rows.size());
+
+        int tsIdx = columns.indexOf("ts");
+        String firstTs = String.valueOf(rows.get(0).get(tsIdx));
+        String lastTs = String.valueOf(rows.get(rows.size() - 1).get(tsIdx));
+        assertTrue("DESC: first " + firstTs + " >= last " + lastTs, firstTs.compareTo(lastTs) >= 0);
+
+        // sub-ms precision survives — pre-fix Stitcher silently coerces to ms (3 digits).
+        assertTrue("expected >3 fractional digits in " + firstTs, firstTs.matches(".*\\.\\d{4,}.*"));
+    }
+
+    private void createParquetBackedIndex() throws Exception {
+        try {
+            client().performRequest(new Request("DELETE", "/" + INDEX));
+        } catch (Exception ignored) {}
+
+        String body = "{"
+            + "\"settings\": {"
+            + "  \"number_of_shards\": " + NUM_SHARDS + ","
+            + "  \"number_of_replicas\": 0,"
+            + "  \"index.pluggable.dataformat.enabled\": true,"
+            + "  \"index.pluggable.dataformat\": \"composite\","
+            + "  \"index.composite.primary_data_format\": \"parquet\","
+            + "  \"index.composite.secondary_data_formats\": \"lucene\""
+            + "},"
+            + "\"mappings\": {"
+            + "  \"properties\": {"
+            + "    \"ts\":       { \"type\": \"date_nanos\" },"
+            + "    \"service\":  { \"type\": \"keyword\"   },"
+            + "    \"severity\": { \"type\": \"keyword\"   },"
+            + "    \"body\":     { \"type\": \"text\", \"store\": true }"
+            + "  }"
+            + "}"
+            + "}";
+
+        Request createIndex = new Request("PUT", "/" + INDEX);
+        createIndex.setJsonEntity(body);
+        Response response = client().performRequest(createIndex);
+        assertOkAndParse(response, "Create index " + INDEX);
+    }
+
+    private void indexDocs() throws Exception {
+        String bulk =
+            "{\"index\":{\"_index\":\"" + INDEX + "\",\"_id\":\"1\"}}\n"
+                + "{\"ts\":\"2025-09-23T00:01:01.123456Z\",\"service\":\"checkout\",\"severity\":\"ERROR\",\"body\":\"failed order: payment declined\"}\n"
+                + "{\"index\":{\"_index\":\"" + INDEX + "\",\"_id\":\"2\"}}\n"
+                + "{\"ts\":\"2025-09-23T00:02:01.234567Z\",\"service\":\"checkout\",\"severity\":\"ERROR\",\"body\":\"failed order due to expired session\"}\n"
+                + "{\"index\":{\"_index\":\"" + INDEX + "\",\"_id\":\"3\"}}\n"
+                + "{\"ts\":\"2025-09-23T00:03:01.345678Z\",\"service\":\"checkout\",\"severity\":\"WARN\",\"body\":\"failed order: inventory check\"}\n"
+                + "{\"index\":{\"_index\":\"" + INDEX + "\",\"_id\":\"4\"}}\n"
+                + "{\"ts\":\"2025-09-23T00:04:01.456789Z\",\"service\":\"checkout\",\"severity\":\"ERROR\",\"body\":\"failed gateway\"}\n"
+                + "{\"index\":{\"_index\":\"" + INDEX + "\",\"_id\":\"5\"}}\n"
+                + "{\"ts\":\"2025-09-23T00:05:00.000000Z\",\"service\":\"frontend\",\"severity\":\"INFO\",\"body\":\"ok\"}\n"
+                + "{\"index\":{\"_index\":\"" + INDEX + "\",\"_id\":\"6\"}}\n"
+                + "{\"ts\":\"2025-09-23T00:06:00.000000Z\",\"service\":\"checkout\",\"severity\":\"INFO\",\"body\":\"successful\"}\n";
+        Request bulkReq = new Request("POST", "/_bulk");
+        bulkReq.addParameter("refresh", "true");
+        bulkReq.setJsonEntity(bulk);
+        assertOkAndParse(client().performRequest(bulkReq), "Bulk index " + INDEX);
+
+        client().performRequest(new Request("POST", "/" + INDEX + "/_flush?force=true"));
+    }
+}