[Analytics Backend] Timechart route end-to-end closure (1/30 → 30/30) (opensearch-project#21747)

Signed-off-by: Kai Huang <huangkaics@gmail.com>

Author: ahkcs

sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ScalarFunction.java

@@ -228,6 +228,24 @@ public enum ScalarFunction {
     DATE_FORMAT(Category.SCALAR, SqlKind.OTHER_FUNCTION),
     TIME_FORMAT(Category.SCALAR, SqlKind.OTHER_FUNCTION),
     STR_TO_DATE(Category.SCALAR, SqlKind.OTHER_FUNCTION),
+    /**
+     * PPL {@code TIMESTAMPDIFF(unit, start, end)} — emits the integer number of {@code unit}s
+     * between two timestamps. Resolves through the SQL plugin's {@code TimestampDiffFunction}
+     * UDF named {@code "TIMESTAMPDIFF"}. Lowering target for PPL `timechart`'s {@code per_*}
+     * aggregations, which expand to {@code DIVIDE(agg * scale, TIMESTAMPDIFF('MILLISECOND',
+     * @timestamp, TIMESTAMPADD(unit, n, @timestamp)))}. The analytics-backend-datafusion adapter
+     * rewrites the call into a DataFusion-native expression because the PPL UDF itself is
+     * unknown to isthmus's default extension catalog.
+     */
+    TIMESTAMPDIFF(Category.SCALAR, SqlKind.OTHER_FUNCTION),
+    /**
+     * PPL {@code TIMESTAMPADD(unit, value, timestamp)} — emits {@code timestamp + value units}.
+     * Resolves through the SQL plugin's {@code TimestampAddFunction} UDF named
+     * {@code "TIMESTAMPADD"}. Lowering target for PPL `timechart`'s {@code per_*}
+     * aggregations (see {@link #TIMESTAMPDIFF}). The analytics-backend-datafusion adapter
+     * rewrites the call into a DataFusion-native interval-add expression.
+     */
+    TIMESTAMPADD(Category.SCALAR, SqlKind.OTHER_FUNCTION),
 
     // ── JSON ────────────────────────────────────────────────────────
     JSON_APPEND(Category.SCALAR, SqlKind.OTHER_FUNCTION),

sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs

@@ -559,11 +559,21 @@ pub unsafe fn sql_to_substrait(
 }
 
 /// Lowers a partial-aggregate Substrait plan against a throwaway session and
-/// returns its physical output schema. NamedTable references are resolved
-/// against empty MemTables built from the substrait base_schema — the plan
-/// itself is the source of truth for the producer side, so no view-type or
-/// timestamp-precision rewrites are applied here. The plan is dropped at
+/// returns its physical output schema **narrowed via
+/// [`schema_coerce::coerce_inferred_schema`]** to types Substrait can bind
+/// against. NamedTable references are resolved against empty MemTables built
+/// from the substrait base_schema — the plan itself is the source of truth for
+/// the producer side, so no view-type or timestamp-precision rewrites are
+/// applied here beyond the post-physical coercer. The plan is dropped at
 /// function exit; only the schema is returned.
+///
+/// <p>The coercer flips Arrow-only types (notably `UInt64` from DataFusion's
+/// `row_number()` physical op) back to their Substrait-compatible counterparts
+/// (`Int64` matching isthmus's `ROW_NUMBER OVER` declaration). The producer
+/// side runs the same coercer + wraps its physical plan with
+/// [`crate::relabel_exec::RelabelExec`] (zero-copy bit-tag flip per mismatched
+/// column), so runtime batches arrive with the same type tags the consumer
+/// registers here — no per-batch cast needed at the partition-stream feed.
 fn derive_schema_from_partial_plan(
     substrait_bytes: &[u8],
 ) -> Result<arrow::datatypes::SchemaRef, DataFusionError> {
@@ -641,7 +651,7 @@ fn derive_schema_from_partial_plan(
 
     let logical_plan = futures::executor::block_on(from_substrait_plan(&session_state, &plan))?;
     let physical_plan = futures::executor::block_on(session_state.create_physical_plan(&logical_plan))?;
-    Ok(physical_plan.schema())
+    Ok(crate::schema_coerce::coerce_inferred_schema(physical_plan.schema()))
 }
 
 /// Encodes a Schema as Arrow IPC stream-format bytes (a schema-only message
@@ -794,6 +804,12 @@ pub unsafe fn register_partition_stream(
     partial_plan_bytes: &[u8],
 ) -> Result<(i64, Vec<u8>), DataFusionError> {
     let session = &mut *(session_ptr as *mut LocalSession);
+    // derive_schema_from_partial_plan applies `schema_coerce::coerce_inferred_schema`
+    // to the physical plan's output schema, matching what isthmus declared on the wire
+    // for the producer (e.g. `Int64` for `ROW_NUMBER OVER`). The producer side runs the
+    // same coercer + wraps its physical plan with `RelabelExec`, so the batches arriving
+    // here through the partition channel are already typed to match this schema — no
+    // per-batch cast at feed time.
     let schema = derive_schema_from_partial_plan(partial_plan_bytes)?;
     let schema_ipc = schema_to_ipc_bytes(schema.as_ref())?;
     let sender = session.register_partition(input_id, schema)?;

sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_executor.rs

@@ -736,6 +736,12 @@ pub async unsafe fn execute_indexed_with_context(
     log_debug!("DataFusion logical plan:\n{}", logical_plan.display_indent());
     let dataframe = ctx.execute_logical_plan(logical_plan).await?;
     let physical_plan = dataframe.create_physical_plan().await?;
+    // Retag bit-compatible Int↔UInt output mismatches to match the substrait-declared
+    // types. The target is schema_coerce::coerce_inferred_schema(physical_schema) — same
+    // narrowing the partition-stream registration uses, so consumer-side StreamingTable
+    // and producer-side batches agree by construction (see crate::relabel_exec).
+    let target_schema = crate::schema_coerce::coerce_inferred_schema(physical_plan.schema());
+    let physical_plan = crate::relabel_exec::wrap_if_relabel_needed(physical_plan, target_schema)?;
     log_debug!("DataFusion physical plan:\n{}", displayable(physical_plan.as_ref()).indent(true));
     let df_stream = execute_stream(physical_plan, ctx.task_ctx())
         .map_err(|e| DataFusionError::Execution(format!("execute_stream: {}", e)))?;

sandbox/plugins/analytics-backend-datafusion/rust/src/lib.rs

@@ -29,6 +29,7 @@ pub mod memory;
 pub mod partition_stream;
 pub mod query_executor;
 pub mod query_tracker;
+pub mod relabel_exec;
 pub mod runtime_manager;
 pub mod schema_coerce;
 pub mod session_context;

sandbox/plugins/analytics-backend-datafusion/rust/src/local_executor.rs

@@ -144,11 +144,12 @@ impl LocalSession {
             DataFusionError::Execution(format!("Failed to decode Substrait plan: {}", e))
         })?;
         let logical_plan = from_substrait_plan(&self.ctx.state(), &plan).await?;
-        self.ctx
-            .execute_logical_plan(logical_plan)
-            .await?
-            .execute_stream()
-            .await
+        let dataframe = self.ctx.execute_logical_plan(logical_plan).await?;
+        let physical_plan = dataframe.create_physical_plan().await?;
+        let target_schema = crate::schema_coerce::coerce_inferred_schema(physical_plan.schema());
+        let physical_plan = crate::relabel_exec::wrap_if_relabel_needed(physical_plan, target_schema)?;
+        datafusion::physical_plan::execute_stream(physical_plan, self.ctx.task_ctx())
+            .map_err(|e| DataFusionError::Execution(format!("execute_substrait: {}", e)))
     }
 
     /// Returns the memory pool the session's `RuntimeEnv` was built with.
@@ -179,6 +180,8 @@ impl LocalSession {
         log_debug!("DataFusion logical plan (reduce):\n{}", logical_plan.display_indent());
         let dataframe = self.ctx.execute_logical_plan(logical_plan).await?;
         let physical_plan = dataframe.create_physical_plan().await?;
+        let target_schema = crate::schema_coerce::coerce_inferred_schema(physical_plan.schema());
+        let physical_plan = crate::relabel_exec::wrap_if_relabel_needed(physical_plan, target_schema)?;
         log_debug!("DataFusion physical plan (reduce):\n{}", displayable(physical_plan.as_ref()).indent(true));
         let stripped = crate::agg_mode::apply_aggregate_mode(
             physical_plan,

sandbox/plugins/analytics-backend-datafusion/rust/src/query_executor.rs

@@ -152,6 +152,13 @@ pub async fn execute_query(
     let logical_plan = from_substrait_plan(&ctx.state(), &substrait_plan).await?;
     let dataframe = ctx.execute_logical_plan(logical_plan).await?;
     let physical_plan = dataframe.create_physical_plan().await?;
+    // Retag any physical-plan output columns whose type tags differ from what Substrait
+    // declared on bit-compatible Int↔UInt pairs (see crate::relabel_exec). The target is
+    // schema_coerce::coerce_inferred_schema(physical_schema) — the same narrowing the
+    // partition-stream registration uses, so the consumer's StreamingTable and the
+    // batches arriving from this producer agree by construction.
+    let target_schema = crate::schema_coerce::coerce_inferred_schema(physical_plan.schema());
+    let physical_plan = crate::relabel_exec::wrap_if_relabel_needed(physical_plan, target_schema)?;
 
     let df_stream = execute_stream(physical_plan, ctx.task_ctx()).map_err(|e| {
         error!("Failed to create execution stream: {}", e);
@@ -197,6 +204,8 @@ pub async fn execute_with_context(
         log_debug!("DataFusion logical plan:\n{}", logical_plan.display_indent());
         let dataframe = handle.ctx.execute_logical_plan(logical_plan).await?;
         let physical_plan = dataframe.create_physical_plan().await?;
+        let target_schema = crate::schema_coerce::coerce_inferred_schema(physical_plan.schema());
+        let physical_plan = crate::relabel_exec::wrap_if_relabel_needed(physical_plan, target_schema)?;
         log_debug!("DataFusion physical plan:\n{}", displayable(physical_plan.as_ref()).indent(true));
 
         let df_stream = execute_stream(physical_plan, handle.ctx.task_ctx()).map_err(|e| {