Introduce DataContext::get(String,Class)

sap-contributions · Oct 9, 2023 · 399ee76 · 399ee76
1 parent eb5c275
commit 399ee76
Show file tree

Hide file tree

Showing 7 changed files with 32 additions and 22 deletions.
diff --git a/core/src/main/java/org/apache/calcite/DataContext.java b/core/src/main/java/org/apache/calcite/DataContext.java
@@ -69,6 +69,19 @@ public interface DataContext {
    */
   @Nullable Object get(String name);
 
+  /**
+   * Returns a context variable.
+   *
+   * <p>Supported variables include: "sparkContext", "currentTimestamp",
+   * "localTimestamp".</p>
+   *
+   * @param name Name of variable
+   * @param clazz Type of variable
+   */
+  @Nullable default <T> T get(String name, Class<T> clazz){
+    return (T) get(name);
+  }
+
   /** Variable that may be asked for in a call to {@link DataContext#get}. */
   enum Variable {
     UTC_TIMESTAMP("utcTimestamp", Long.class),

diff --git a/core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableLimit.java b/core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableLimit.java
@@ -121,11 +121,9 @@ public static EnumerableLimit create(final RelNode input, @Nullable RexNode offs
   static Expression getExpression(RexNode rexNode) {
     if (rexNode instanceof RexDynamicParam) {
       final RexDynamicParam param = (RexDynamicParam) rexNode;
-      return Expressions.convert_(
-          Expressions.call(DataContext.ROOT,
-              BuiltInMethod.DATA_CONTEXT_GET.method,
-              Expressions.constant("?" + param.getIndex())),
-          Integer.class);
+      return Expressions.call(DataContext.ROOT,
+              BuiltInMethod.DATA_CONTEXT_GET_TYPED.method,
+              Expressions.constant("?" + param.getIndex()),Expressions.constant(Integer.class));
     } else {
       return Expressions.constant(RexLiteral.intValue(rexNode));
     }

diff --git a/core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableRelImplementor.java b/core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableRelImplementor.java
@@ -157,11 +157,10 @@ public ClassDeclaration implementRoot(EnumerableRel rootRel,
     final Collection<Statement> stashed =
         Collections2.transform(stashedParameters.values(),
             input -> Expressions.declare(Modifier.FINAL, input,
-                Expressions.convert_(
-                    Expressions.call(DataContext.ROOT,
-                        BuiltInMethod.DATA_CONTEXT_GET.method,
-                        Expressions.constant(input.name)),
-                    input.type)));
+                Expressions.call(DataContext.ROOT,
+                    BuiltInMethod.DATA_CONTEXT_GET_TYPED.method,
+                    Expressions.constant(input.name),
+                    Expressions.constant(input.type))));
 
     final BlockStatement block =
         Expressions.block(

diff --git a/core/src/main/java/org/apache/calcite/adapter/enumerable/RexToLixTranslator.java b/core/src/main/java/org/apache/calcite/adapter/enumerable/RexToLixTranslator.java
@@ -1357,8 +1357,9 @@ private Result toInnerStorageType(Result result, Type storageType) {
         ? currentStorageType : typeFactory.getJavaClass(dynamicParam.getType());
     final Expression valueExpression =
         EnumUtils.convert(
-            Expressions.call(root, BuiltInMethod.DATA_CONTEXT_GET.method,
-                Expressions.constant("?" + dynamicParam.getIndex())),
+            Expressions.call(root, BuiltInMethod.DATA_CONTEXT_GET_TYPED.method,
+                Expressions.constant("?" + dynamicParam.getIndex()),
+                Expressions.constant(storageType)),
             storageType);
     final ParameterExpression valueVariable =
         Expressions.parameter(valueExpression.getType(),

diff --git a/core/src/main/java/org/apache/calcite/sql/advise/SqlAdvisorGetHintsFunction.java b/core/src/main/java/org/apache/calcite/sql/advise/SqlAdvisorGetHintsFunction.java
@@ -50,11 +50,10 @@
 public class SqlAdvisorGetHintsFunction
     implements TableFunction, ImplementableFunction {
   private static final Expression ADVISOR =
-      Expressions.convert_(
-          Expressions.call(DataContext.ROOT,
-              BuiltInMethod.DATA_CONTEXT_GET.method,
-              Expressions.constant(DataContext.Variable.SQL_ADVISOR.camelName)),
-          SqlAdvisor.class);
+      Expressions.call(DataContext.ROOT,
+          BuiltInMethod.DATA_CONTEXT_GET_TYPED.method,
+          Expressions.constant(DataContext.Variable.SQL_ADVISOR.camelName),
+          Expressions.constant(SqlAdvisor.class));
 
   private static final Method GET_COMPLETION_HINTS =
       Types.lookupMethod(SqlAdvisorGetHintsFunction.class, "getCompletionHints",

diff --git a/core/src/main/java/org/apache/calcite/sql/advise/SqlAdvisorGetHintsFunction2.java b/core/src/main/java/org/apache/calcite/sql/advise/SqlAdvisorGetHintsFunction2.java
@@ -52,11 +52,10 @@
 public class SqlAdvisorGetHintsFunction2
     implements TableFunction, ImplementableFunction {
   private static final Expression ADVISOR =
-      Expressions.convert_(
-          Expressions.call(DataContext.ROOT,
-              BuiltInMethod.DATA_CONTEXT_GET.method,
-              Expressions.constant(DataContext.Variable.SQL_ADVISOR.camelName)),
-          SqlAdvisor.class);
+      Expressions.call(DataContext.ROOT,
+          BuiltInMethod.DATA_CONTEXT_GET_TYPED.method,
+          Expressions.constant(DataContext.Variable.SQL_ADVISOR.camelName),
+          Expressions.constant(SqlAdvisor.class));
 
   private static final Method GET_COMPLETION_HINTS =
       Types.lookupMethod(SqlAdvisorGetHintsFunction2.class, "getCompletionHints",

diff --git a/core/src/main/java/org/apache/calcite/util/BuiltInMethod.java b/core/src/main/java/org/apache/calcite/util/BuiltInMethod.java
@@ -170,6 +170,7 @@ public enum BuiltInMethod {
       SchemaPlus.class, Class.class, String.class),
   REFLECTIVE_SCHEMA_GET_TARGET(ReflectiveSchema.class, "getTarget"),
   DATA_CONTEXT_GET(DataContext.class, "get", String.class),
+  DATA_CONTEXT_GET_TYPED(DataContext.class, "get", String.class, Class.class),
   DATA_CONTEXT_GET_ROOT_SCHEMA(DataContext.class, "getRootSchema"),
   JDBC_SCHEMA_DATA_SOURCE(JdbcSchema.class, "getDataSource"),
   ROW_VALUE(Row.class, "getObject", int.class),