opensearch-project · mch2 · May 6, 2026 · May 5, 2026 · May 6, 2026 · May 6, 2026
@@ -0,0 +1,96 @@
+/*
+ * 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.spi;
+
+import org.apache.calcite.plan.RelOptCluster;
+import org.apache.calcite.rex.RexBuilder;
+import org.apache.calcite.rex.RexCall;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.sql.SqlOperator;
+import org.apache.calcite.sql.type.SqlTypeName;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * Reusable base for {@link ScalarFunctionAdapter}s that rewrite a Calcite call
+ * to a different named target, optionally prepending or appending literal
+ * operands. Pure shape rewriting — no decomposition into a different semantic
+ * function. For that use case (e.g. {@code ILIKE → LIKE(LOWER(a), LOWER(b))})
+ * write a dedicated adapter instead.
+ *
+ * <p>Example use:
+ * <pre>
+ *   class YearAdapter extends AbstractNameMappingAdapter {
+ *       YearAdapter() {
+ *           super(SqlLibraryOperators.DATE_PART, List.of("year"), List.of());
+ *       }
+ *   }
+ * </pre>
+ * rewrites {@code YEAR(ts)} to {@code date_part('year', ts)}. Paired with the
+ * {@code date_part} signature in a backend's extension catalog so the isthmus
+ * visitor resolves it against the backend's native date_part.
+ *
+ * @opensearch.internal
+ */
+public abstract class AbstractNameMappingAdapter implements ScalarFunctionAdapter {
+
+    private final SqlOperator targetOperator;
+    private final List<Object> prependLiterals;
+    private final List<Object> appendLiterals;
+
+    /**
+     * @param targetOperator  the Calcite {@link SqlOperator} the rewritten call
+     *                        will use. The isthmus visitor resolves this to a
+     *                        Substrait invocation against the backend's loaded
+     *                        extension catalog.
+     * @param prependLiterals literals to prepend to the operand list (e.g.
+     *                        {@code List.of("year")} to prepend a string literal).
+     *                        Currently supports {@link String}, {@link Integer},
+     *                        {@link Long}, {@link Double}, {@link Boolean}.
+     * @param appendLiterals  literals to append to the operand list.
+     */
+    protected AbstractNameMappingAdapter(SqlOperator targetOperator, List<Object> prependLiterals, List<Object> appendLiterals) {
+        this.targetOperator = targetOperator;
+        this.prependLiterals = List.copyOf(prependLiterals);
+        this.appendLiterals = List.copyOf(appendLiterals);
+    }
+
+    @Override
+    public RexNode adapt(RexCall original, List<FieldStorageInfo> fieldStorage, RelOptCluster cluster) {
+        RexBuilder rexBuilder = cluster.getRexBuilder();
+        List<RexNode> operands = new ArrayList<>(original.getOperands().size() + prependLiterals.size() + appendLiterals.size());
+        for (Object literal : prependLiterals) {
+            operands.add(rexBuilder.makeLiteral(literal, inferLiteralType(rexBuilder, literal), true));
+        }
+        operands.addAll(original.getOperands());
+        for (Object literal : appendLiterals) {
+            operands.add(rexBuilder.makeLiteral(literal, inferLiteralType(rexBuilder, literal), true));
+        }
+        // Preserve the original call's return type. The enclosing operator (Project
+        // / Filter) caches its rowType from the pre-adaptation expression; if the
+        // rewritten call's Calcite-inferred type differs (e.g. PPL YEAR returns
+        // INTEGER but SqlLibraryOperators.DATE_PART is SqlExtractFunction → BIGINT),
+        // the downstream stripAnnotations path feeds the adapted expr into
+        // LogicalProject.create together with the cached rowType, and
+        // Project.isValid's compatibleTypes check throws an AssertionError that
+        // breaks fragment conversion.
+        return rexBuilder.makeCall(original.getType(), targetOperator, operands);
+    }
+
+    private static org.apache.calcite.rel.type.RelDataType inferLiteralType(RexBuilder rexBuilder, Object literal) {
+        var typeFactory = rexBuilder.getTypeFactory();
+        if (literal instanceof String) return typeFactory.createSqlType(SqlTypeName.VARCHAR);
+        if (literal instanceof Integer) return typeFactory.createSqlType(SqlTypeName.INTEGER);
+        if (literal instanceof Long) return typeFactory.createSqlType(SqlTypeName.BIGINT);
+        if (literal instanceof Double) return typeFactory.createSqlType(SqlTypeName.DOUBLE);
+        if (literal instanceof Boolean) return typeFactory.createSqlType(SqlTypeName.BOOLEAN);
+        throw new IllegalArgumentException("Unsupported literal type: " + literal.getClass());
+    }
+}
@@ -77,7 +77,10 @@ public enum ScalarFunction {
     EXTRACT(Category.SCALAR, SqlKind.EXTRACT),
 
     // ── Datetime ────────────────────────────────────────────────────
-    TIMESTAMP(Category.SCALAR, SqlKind.OTHER_FUNCTION);
+    TIMESTAMP(Category.SCALAR, SqlKind.OTHER_FUNCTION),
+    YEAR(Category.SCALAR, SqlKind.OTHER_FUNCTION),
+    CONVERT_TZ(Category.SCALAR, SqlKind.OTHER_FUNCTION),
+    UNIX_TIMESTAMP(Category.SCALAR, SqlKind.OTHER_FUNCTION);
 
     /**
      * Category of scalar function.
@@ -87,7 +90,9 @@ public enum Category {
         FULL_TEXT,
         STRING,
         MATH,
-        /** Catch-all for functions that don't fit other categories (CAST, CASE, COALESCE, EXTRACT, etc.). */
+        /**
+         * Catch-all for functions that don't fit other categories (CAST, CASE, COALESCE, EXTRACT, etc.).
+         */
         SCALAR
     }
 
@@ -121,7 +126,9 @@ public static ScalarFunction fromSqlKind(SqlKind kind) {
         return null;
     }
 
-    /** Maps a Calcite SqlFunction to a ScalarFunction by name, or throws if not recognized. */
+    /**
+     * Maps a Calcite SqlFunction to a ScalarFunction by name, or null if not recognized.
+     */
     public static ScalarFunction fromSqlFunction(SqlFunction function) {
         // TODO: Add an explicit functionName field per enum constant instead of relying on
         // valueOf(toUpperCase). This couples enum constant naming to SQL function naming convention.

@@ -43,6 +43,9 @@ chrono = { workspace = true }
 roaring = "0.10"
 thiserror = { workspace = true }
 
+# convert_tz UDF
+chrono-tz = "0.10"
+
 [dev-dependencies]
 criterion = { workspace = true }
 tempfile = { workspace = true }

@@ -443,6 +443,7 @@ pub unsafe fn sql_to_substrait(
             .with_default_features()
             .build();
         let ctx = datafusion::prelude::SessionContext::new_with_state(state);
+        crate::udf::register_all(&ctx);
 
         let listing_options = ListingOptions::new(Arc::new(ParquetFormat::new()))
             .with_file_extension(".parquet")

@@ -133,6 +133,7 @@ pub async fn execute_indexed_query(
         .build();
     let ctx = SessionContext::new_with_state(state);
     ctx.register_udf(create_index_filter_udf());
+    crate::udf::register_all(&ctx);
 
     // Resolve the object store for this shard's table URL (file://, s3://,
     // gs://, ... whatever the global runtime has registered). We pass this

@@ -30,3 +30,4 @@ pub mod query_memory_pool_tracker;
 pub mod runtime_manager;
 pub mod session_context;
 pub mod statistics_cache;
+pub mod udf;
@@ -70,9 +70,9 @@ impl LocalSession {
             .with_runtime_env(runtime_env)
             .with_default_features()
             .build();
-        Self {
-            ctx: SessionContext::new_with_state(state),
-        }
+        let ctx = SessionContext::new_with_state(state);
+        crate::udf::register_all(&ctx);
+        Self { ctx }
     }
 
     /// Registers a streaming input on the session under `name` and returns the

@@ -100,6 +100,7 @@ pub async fn execute_query(
         .build();
 
     let ctx = SessionContext::new_with_state(state);
+    crate::udf::register_all(&ctx);
 
     // Register table via ListingTable — all IO goes through object store
     let file_format = ParquetFormat::new();

@@ -99,6 +99,7 @@ pub async unsafe fn create_session_context(
         .build();
 
     let ctx = SessionContext::new_with_state(state);
+    crate::udf::register_all(&ctx);
 
     // Register default ListingTable for parquet scans
     let listing_options = ListingOptions::new(Arc::new(ParquetFormat::new()))