SQL: Preserve original source for cast/convert function (#40271)

Improve rule for pruning cast to preserve the original source
Fix #40239

(cherry picked from commit 7591cb1a1577320b3aec2ec557b0f881b6af744f)

x-pack/plugin/sql/qa/src/main/resources/agg.csv-spec

@@ -85,8 +85,8 @@ SELECT SUM(salary) FROM test_emp;
 aggregateWithCastPruned
 SELECT CAST(SUM(salary) AS INTEGER) FROM test_emp;
 
-  SUM(salary)
+ CAST(SUM(salary) AS INTEGER)
--------------
+-----------------------------
 4824855
 ;
 

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

@@ -163,6 +163,7 @@ public class Optimizer extends RuleExecutor<LogicalPlan> {
                 );
         //new BalanceBooleanTrees());
         Batch label = new Batch("Set as Optimized", Limiter.ONCE,
+                CleanAliases.INSTANCE,
                 new SetAsOptimized());
 
         return Arrays.asList(operators, aggregate, local, label);
@@ -895,7 +896,7 @@ public class Optimizer extends RuleExecutor<LogicalPlan> {
                                     // Check if the groupings (a, y) match the orderings (b, x) through the aggregates' aliases (x, y)
                                     // e.g. SELECT a AS x, b AS y ... GROUP BY a, y ORDER BY b, x
                                     if ((equalsAsAttribute(child, group)
-                                            && (equalsAsAttribute(alias, fieldToOrder) || equalsAsAttribute(child, fieldToOrder))) 
+                                            && (equalsAsAttribute(alias, fieldToOrder) || equalsAsAttribute(child, fieldToOrder)))
                                         || (equalsAsAttribute(alias, group)
                                                 && (equalsAsAttribute(alias, fieldToOrder) || equalsAsAttribute(child, fieldToOrder)))) {
                                         isMatching.set(Boolean.TRUE);
@@ -944,43 +945,22 @@ public class Optimizer extends RuleExecutor<LogicalPlan> {
         protected LogicalPlan rule(LogicalPlan plan) {
             final Map<Attribute, Attribute> replacedCast = new LinkedHashMap<>();
 
-            // first eliminate casts inside Aliases
+            // eliminate redundant casts
             LogicalPlan transformed = plan.transformExpressionsUp(e -> {
-                // cast wrapped in an alias
-                if (e instanceof Alias) {
-                    Alias as = (Alias) e;
-                    if (as.child() instanceof Cast) {
-                        Cast c = (Cast) as.child();
-
-                        if (c.from() == c.to()) {
-                            Alias newAs = new Alias(as.source(), as.name(), as.qualifier(), c.field(), as.id(), as.synthetic());
-                            replacedCast.put(as.toAttribute(), newAs.toAttribute());
-                            return newAs;
-                        }
-                    }
-                    return e;
-                }
-                return e;
-            });
-
-            // then handle stand-alone casts (mixed together the cast rule will kick in before the alias)
-            transformed = transformed.transformExpressionsUp(e -> {
                 if (e instanceof Cast) {
                     Cast c = (Cast) e;
 
                     if (c.from() == c.to()) {
                         Expression argument = c.field();
-                        if (argument instanceof NamedExpression) {
+                        Alias as = new Alias(c.source(), c.sourceText(), argument);
-                            replacedCast.put(c.toAttribute(), ((NamedExpression) argument).toAttribute());
+                        replacedCast.put(c.toAttribute(), as.toAttribute());
-                        }
 
-                        return argument;
+                        return as;
                     }
                 }
                 return e;
             });
 
-
             // replace attributes from previous removed Casts
             if (!replacedCast.isEmpty()) {
                 return transformed.transformUp(p -> {

x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/parser/SqlParserTests.java

@@ -13,6 +13,7 @@ import org.elasticsearch.xpack.sql.expression.Order;
 import org.elasticsearch.xpack.sql.expression.UnresolvedAttribute;
 import org.elasticsearch.xpack.sql.expression.UnresolvedStar;
 import org.elasticsearch.xpack.sql.expression.function.UnresolvedFunction;
+import org.elasticsearch.xpack.sql.expression.function.scalar.Cast;
 import org.elasticsearch.xpack.sql.expression.predicate.fulltext.MatchQueryPredicate;
 import org.elasticsearch.xpack.sql.expression.predicate.fulltext.MultiMatchQueryPredicate;
 import org.elasticsearch.xpack.sql.expression.predicate.fulltext.StringQueryPredicate;
@@ -65,6 +66,21 @@ public class SqlParserTests extends ESTestCase {
         assertEquals("SCORE()", f.sourceText());
     }
 
+    public void testSelectCast() {
+        Cast f = singleProjection(project(parseStatement("SELECT CAST(POWER(languages, 2) AS DOUBLE) FROM foo")), Cast.class);
+        assertEquals("CAST(POWER(languages, 2) AS DOUBLE)", f.sourceText());
+    }
+
+    public void testSelectCastOperator() {
+        Cast f = singleProjection(project(parseStatement("SELECT POWER(languages, 2)::DOUBLE FROM foo")), Cast.class);
+        assertEquals("POWER(languages, 2)::DOUBLE", f.sourceText());
+    }
+
+    public void testSelectCastWithSQLOperator() {
+        Cast f = singleProjection(project(parseStatement("SELECT CONVERT(POWER(languages, 2), SQL_DOUBLE) FROM foo")), Cast.class);
+        assertEquals("CONVERT(POWER(languages, 2), SQL_DOUBLE)", f.sourceText());
+    }
+
     public void testSelectAddWithParanthesis() {
         Add f = singleProjection(project(parseStatement("SELECT (1 +  2)")), Add.class);
         assertEquals("1 +  2", f.sourceText());