SQL: Make Literal a NamedExpression (#33583)

* SQL: Make Literal a NamedExpression

Literal now is a NamedExpression reducing the need for Aliases for
folded expressions leading to simpler optimization rules.

Fix #33523

(cherry picked from commit 91bca17)

Author: costin

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/expression/Attribute.java

@@ -7,30 +7,28 @@
 
 import org.elasticsearch.xpack.sql.tree.Location;
 
+import java.util.List;
 import java.util.Objects;
 
 import static java.util.Collections.emptyList;
 
-import java.util.List;
-
 /**
- * {@link Expression}s that can be converted into Elasticsearch
- * sorts, aggregations, or queries. They can also be extracted
- * from the result of a search.
+ * {@link Expression}s that can be materialized and represent the result columns sent to the client.
+ * Typically are converted into constants, functions or Elasticsearch order-bys,
+ * aggregations, or queries. They can also be extracted from the result of a search.
  *
  * In the statement {@code SELECT ABS(foo), A, B+C FROM ...} the three named
- * expressions (ABS(foo), A, B+C) get converted to attributes and the user can
+ * expressions {@code ABS(foo), A, B+C} get converted to attributes and the user can
  * only see Attributes.
  *
- * In the statement {@code SELECT foo FROM TABLE WHERE foo > 10 + 1} 10+1 is an
- * expression. It's not named - meaning there's no alias for it (defined by the
- * user) and as such there's no attribute - no column to be returned to the user.
- * It's an expression used for filtering so it doesn't appear in the result set
- * (derived table). "foo" on the other hand is an expression, a named expression
- * (it has a name) and also an attribute - it's a column in the result set.
+ * In the statement {@code SELECT foo FROM TABLE WHERE foo > 10 + 1} both {@code foo} and
+ * {@code 10 + 1} are named expressions, the first due to the SELECT, the second due to being a function.
+ * However since {@code 10 + 1} is used for filtering it doesn't appear appear in the result set
+ * (derived table) and as such it is never translated to an attribute.
+ * "foo" on the other hand is since it's a column in the result set.
  *
- * Another example {@code SELECT foo FROM ... WHERE bar > 10 +1} "foo" gets
- * converted into an Attribute, bar does not. That's because bar is used for
+ * Another example {@code SELECT foo FROM ... WHERE bar > 10 +1} {@code foo} gets
+ * converted into an Attribute, bar does not. That's because {@code bar} is used for
  * filtering alone but it's not part of the projection meaning the user doesn't
  * need it in the derived table.
  */

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/expression/Expressions.java

@@ -82,13 +82,7 @@ public static AttributeSet references(List<? extends Expression> exps) {
     }
 
     public static String name(Expression e) {
-        if (e instanceof NamedExpression) {
-            return ((NamedExpression) e).name();
-        } else if (e instanceof Literal) {
-            return e.toString();
-        } else {
-            return e.nodeName();
-        }
+        return e instanceof NamedExpression ? ((NamedExpression) e).name() : e.nodeName();
     }
 
     public static List<String> names(Collection<? extends Expression> e) {
@@ -105,7 +99,7 @@ public static Attribute attribute(Expression e) {
             return ((NamedExpression) e).toAttribute();
         }
         if (e != null && e.foldable()) {
-            return new LiteralAttribute(Literal.of(e));
+            return Literal.of(e).toAttribute();
         }
         return null;
     }

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/expression/Literal.java

@@ -12,17 +12,28 @@
 import org.elasticsearch.xpack.sql.type.DataTypeConversion;
 import org.elasticsearch.xpack.sql.type.DataTypes;
 
+import java.util.List;
 import java.util.Objects;
 
-public class Literal extends LeafExpression {
+import static java.util.Collections.emptyList;
+
+/**
+ * SQL Literal or constant.
+ */
+public class Literal extends NamedExpression {
+
     public static final Literal TRUE = Literal.of(Location.EMPTY, Boolean.TRUE);
     public static final Literal FALSE = Literal.of(Location.EMPTY, Boolean.FALSE);
 
     private final Object value;
     private final DataType dataType;
 
     public Literal(Location location, Object value, DataType dataType) {
-        super(location);
+        this(location, null, value, dataType);
+    }
+
+    public Literal(Location location, String name, Object value, DataType dataType) {
+        super(location, name == null ? String.valueOf(value) : name, emptyList(), null);
         this.dataType = dataType;
         this.value = DataTypeConversion.convert(value, dataType);
     }
@@ -61,48 +72,83 @@ public Object fold() {
         return value;
     }
 
+    @Override
+    public Attribute toAttribute() {
+        return new LiteralAttribute(location(), name(), null, false, id(), false, dataType, this);
+    }
+
+    @Override
+    public Expression replaceChildren(List<Expression> newChildren) {
+        throw new UnsupportedOperationException("this type of node doesn't have any children to replace");
+    }
+
+    @Override
+    public AttributeSet references() {
+        return AttributeSet.EMPTY;
+    }
 
     @Override
     public int hashCode() {
-        return Objects.hash(value, dataType);
+        return Objects.hash(name(), value, dataType);
     }
 
     @Override
     public boolean equals(Object obj) {
         if (this == obj) {
             return true;
         }
-
         if (obj == null || getClass() != obj.getClass()) {
             return false;
         }
 
         Literal other = (Literal) obj;
-        return Objects.equals(value, other.value)
+        return Objects.equals(name(), other.name())
+                && Objects.equals(value, other.value)
                 && Objects.equals(dataType, other.dataType);
     }
 
     @Override
     public String toString() {
-        return Objects.toString(value);
+        String s = String.valueOf(value);
+        return name().equals(s) ? s : name() + "=" + value;
     }
 
+    /**
+     * Utility method for creating 'in-line' Literals (out of values instead of expressions).
+     */
     public static Literal of(Location loc, Object value) {
         if (value instanceof Literal) {
             return (Literal) value;
         }
         return new Literal(loc, value, DataTypes.fromJava(value));
     }
 
+    /**
+     * Utility method for creating a literal out of a foldable expression.
+     * Throws an exception if the expression is not foldable.
+     */
     public static Literal of(Expression foldable) {
-        if (foldable instanceof Literal) {
-            return (Literal) foldable;
-        }
+        return of((String) null, foldable);
+    }
 
+    public static Literal of(String name, Expression foldable) {
         if (!foldable.foldable()) {
             throw new SqlIllegalArgumentException("Foldable expression required for Literal creation; received unfoldable " + foldable);
         }
 
-        return new Literal(foldable.location(), foldable.fold(), foldable.dataType());
+        if (foldable instanceof Literal) {
+            Literal l = (Literal) foldable;
+            if (name == null || l.name().equals(name)) {
+                return l;
+            }
+        }
+
+        Object fold = foldable.fold();
+
+        if (name == null) {
+            name = foldable instanceof NamedExpression ? ((NamedExpression) foldable).name() : String.valueOf(fold);
+        }
+
+        return new Literal(foldable.location(), name, fold, foldable.dataType());
     }
-}
+}

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/expression/LiteralAttribute.java

@@ -15,20 +15,12 @@ public class LiteralAttribute extends TypedAttribute {
 
     private final Literal literal;
 
-    public LiteralAttribute(Literal literal) {
-        this(literal.location(), String.valueOf(literal.fold()), null, false, null, false, literal.dataType(), literal);
-    }
-
     public LiteralAttribute(Location location, String name, String qualifier, boolean nullable, ExpressionId id, boolean synthetic,
             DataType dataType, Literal literal) {
         super(location, name, dataType, qualifier, nullable, id, synthetic);
         this.literal = literal;
     }
 
-    public Literal literal() {
-        return literal;
-    }
-
     @Override
     protected NodeInfo<LiteralAttribute> info() {
         return NodeInfo.create(this, LiteralAttribute::new,
@@ -49,4 +41,4 @@ public ProcessorDefinition asProcessorDefinition() {
     protected String label() {
         return "c";
     }
-}
+}

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/expression/function/scalar/ScalarFunction.java

@@ -10,7 +10,6 @@
 import org.elasticsearch.xpack.sql.expression.Expression;
 import org.elasticsearch.xpack.sql.expression.Expressions;
 import org.elasticsearch.xpack.sql.expression.FieldAttribute;
-import org.elasticsearch.xpack.sql.expression.LiteralAttribute;
 import org.elasticsearch.xpack.sql.expression.function.Function;
 import org.elasticsearch.xpack.sql.expression.function.aggregate.AggregateFunctionAttribute;
 import org.elasticsearch.xpack.sql.expression.function.scalar.processor.definition.ProcessorDefinition;
@@ -69,11 +68,9 @@ protected ScriptTemplate asScript(Expression exp) {
             if (attr instanceof AggregateFunctionAttribute) {
                 return asScriptFrom((AggregateFunctionAttribute) attr);
             }
-            if (attr instanceof LiteralAttribute) {
-                return asScriptFrom((LiteralAttribute) attr);
+            if (attr instanceof FieldAttribute) {
+                return asScriptFrom((FieldAttribute) attr);
             }
-            // fall-back to
-            return asScriptFrom((FieldAttribute) attr);
         }
         throw new SqlIllegalArgumentException("Cannot evaluate script for expression {}", exp);
     }
@@ -102,12 +99,6 @@ protected ScriptTemplate asScriptFrom(AggregateFunctionAttribute aggregate) {
                 aggregate.dataType());
     }
 
-    protected ScriptTemplate asScriptFrom(LiteralAttribute literal) {
-        return new ScriptTemplate(formatScript("{}"),
-            paramsBuilder().variable(literal.literal()).build(),
-            literal.dataType());
-    }
-
     protected String formatScript(String scriptTemplate) {
         return formatTemplate(scriptTemplate);
     }

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/expression/function/scalar/math/E.java

@@ -21,7 +21,7 @@ public class E extends MathFunction {
     private static final ScriptTemplate TEMPLATE = new ScriptTemplate("Math.E", Params.EMPTY, DataType.DOUBLE);
 
     public E(Location location) {
-        super(location, new Literal(location, Math.E, DataType.DOUBLE));
+        super(location, new Literal(location, "E", Math.E, DataType.DOUBLE));
     }
 
     @Override

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/expression/function/scalar/math/Pi.java

@@ -21,7 +21,7 @@ public class Pi extends MathFunction {
     private static final ScriptTemplate TEMPLATE = new ScriptTemplate("Math.PI", Params.EMPTY, DataType.DOUBLE);
 
     public Pi(Location location) {
-        super(location, new Literal(location, Math.PI, DataType.DOUBLE));
+        super(location, new Literal(location, "PI", Math.PI, DataType.DOUBLE));
     }
 
     @Override

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/optimizer/Optimizer.java

@@ -1118,36 +1118,12 @@ static class ConstantFolding extends OptimizerExpressionRule {
 
         @Override
         protected Expression rule(Expression e) {
-            // handle aliases to avoid double aliasing of functions
-            // alias points to function which gets folded and wrapped in an alias that is
-            // aliases
             if (e instanceof Alias) {
                 Alias a = (Alias) e;
-                Expression fold = fold(a.child());
-                if (fold != a.child()) {
-                    return new Alias(a.location(), a.name(), null, fold, a.id());
-                }
-                return a;
-            }
-
-            Expression fold = fold(e);
-            if (fold != e) {
-                // preserve the name through an alias
-                if (e instanceof NamedExpression) {
-                    NamedExpression ne = (NamedExpression) e;
-                    return new Alias(e.location(), ne.name(), null, fold, ne.id());
-                }
-                return fold;
+                return a.child().foldable() ? Literal.of(a.name(), a.child()) : a;
             }
-            return e;
-        }
 
-        private Expression fold(Expression e) {
-            // literals are always foldable, so avoid creating a duplicate
-            if (e.foldable() && !(e instanceof Literal)) {
-                return new Literal(e.location(), e.fold(), e.dataType());
-            }
-            return e;
+            return e.foldable() ? Literal.of(e) : e;
         }
     }
 
@@ -1836,14 +1812,11 @@ protected LogicalPlan rule(LogicalPlan plan) {
         private List<Object> extractConstants(List<? extends NamedExpression> named) {
             List<Object> values = new ArrayList<>();
             for (NamedExpression n : named) {
-                if (n instanceof Alias) {
-                    Alias a = (Alias) n;
-                    if (a.child().foldable()) {
-                        values.add(a.child().fold());
-                    }
-                    else {
-                        return values;
-                    }
+                if (n.foldable()) {
+                    values.add(n.fold());
+                } else {
+                    // not everything is foldable, bail-out early
+                    return values;
                 }
             }
             return values;

x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/expression/LiteralTests.java

@@ -61,7 +61,7 @@ protected Literal randomInstance() {
 
     @Override
     protected Literal copy(Literal instance) {
-        return new Literal(instance.location(), instance.value(), instance.dataType());
+        return new Literal(instance.location(), instance.name(), instance.value(), instance.dataType());
     }
 
     @Override