SQL: fix scripting for grouped by datetime functions

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

@@ -309,3 +309,97 @@ SELECT birth_date, MAX(hire_date) - INTERVAL 1 YEAR AS f FROM test_emp GROUP BY
 1952-06-13T00:00:00Z|1953           
 1952-07-08T00:00:00Z|1953   
 ;
+
+monthOfDatePlusInterval_And_GroupBy
+SELECT WEEK_OF_YEAR(birth_date + INTERVAL 25 YEAR) x, COUNT(*) c FROM test_emp GROUP BY x HAVING c >= 3 ORDER BY c DESC;
+
+      x:i      |       c:l       
+---------------+---------------
+null           |10             
+22             |6              
+4              |4              
+16             |4              
+30             |4              
+40             |4              
+45             |4              
+1              |3              
+8              |3              
+21             |3              
+28             |3              
+32             |3              
+37             |3
+;
+
+dayOfWeekPlusInterval_And_GroupBy
+SELECT DOW(birth_date + INTERVAL 5 YEAR) x, COUNT(*) c FROM test_emp GROUP BY x HAVING c >= 3 ORDER BY c DESC;
+
+      x:i      |       c:l       
+---------------+---------------
+2              |18             
+3              |15             
+5              |15             
+4              |12             
+6              |12             
+null           |10             
+7              |10             
+1              |8              
+;
+
+dayNamePlusInterval_And_GroupBy
+SELECT DAY_NAME(birth_date + INTERVAL 5 YEAR) x, COUNT(*) c FROM test_emp GROUP BY x HAVING c >= 10 ORDER BY c DESC;
+
+      x:s      |       c:l       
+---------------+---------------
+Monday         |18             
+Thursday       |15             
+Tuesday        |15             
+Friday         |12             
+Wednesday      |12             
+null           |10             
+Saturday       |10          
+;
+
+monthNamePlusInterval_And_GroupBy
+SELECT MONTH_NAME(birth_date + INTERVAL 5 YEAR) x, COUNT(*) c FROM test_emp GROUP BY x HAVING c >= 5 ORDER BY c DESC;
+
+      x:s      |       c:l       
+---------------+---------------
+null           |10             
+May            |10             
+September      |10             
+July           |9              
+October        |9              
+April          |8              
+February       |8              
+November       |8              
+December       |7              
+June           |7              
+August         |6              
+January        |6             
+;
+
+quarterPlusInterval_And_GroupBy
+SELECT QUARTER(birth_date + INTERVAL 5 YEAR) x, COUNT(*) c FROM test_emp GROUP BY x HAVING c >= 5 ORDER BY x DESC;
+
+      x:i      |       c:l       
+---------------+---------------
+4              |24             
+3              |25             
+2              |25             
+1              |16             
+null           |10             
+;
+
+dayOfMonthPlusInterval_And_GroupBy
+SELECT DOM(birth_date + INTERVAL 5 YEAR) x, COUNT(*) c FROM test_emp GROUP BY x HAVING c >= 5 ORDER BY x DESC;
+
+      x:i      |       c:l       
+---------------+---------------
+25             |5              
+23             |6              
+21             |5              
+19             |7              
+7              |5              
+1              |5              
+null           |10             
+;

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

@@ -49,7 +49,6 @@ public ScriptTemplate asScript() {
               .variable(extractor.chronoField().name());
 
         return new ScriptTemplate(template, params.build(), dataType());
-
     }
 
     @Override

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

@@ -6,18 +6,16 @@
 package org.elasticsearch.xpack.sql.expression.function.scalar.datetime;
 
 import org.elasticsearch.xpack.sql.expression.Expression;
-import org.elasticsearch.xpack.sql.expression.FieldAttribute;
 import org.elasticsearch.xpack.sql.expression.function.scalar.datetime.NamedDateTimeProcessor.NameExtractor;
 import org.elasticsearch.xpack.sql.expression.gen.processor.Processor;
+import org.elasticsearch.xpack.sql.expression.gen.script.ParamsBuilder;
 import org.elasticsearch.xpack.sql.expression.gen.script.ScriptTemplate;
 import org.elasticsearch.xpack.sql.tree.Source;
 import org.elasticsearch.xpack.sql.type.DataType;
 import org.elasticsearch.xpack.sql.util.StringUtils;
 
 import java.time.ZoneId;
-import java.util.Locale;
 
-import static java.lang.String.format;
 import static org.elasticsearch.xpack.sql.expression.gen.script.ParamsBuilder.paramsBuilder;
 
 /*
@@ -33,14 +31,16 @@ abstract class NamedDateTimeFunction extends BaseDateTimeFunction {
     }
 
     @Override
-    public ScriptTemplate scriptWithField(FieldAttribute field) {
-        return new ScriptTemplate(
-                formatTemplate(format(Locale.ROOT, "{sql}.%s(doc[{}].value, {})",
-                        StringUtils.underscoreToLowerCamelCase(nameExtractor.name()))),
-                paramsBuilder()
-                  .variable(field.name())
-                  .variable(zoneId().getId()).build(),
-                dataType());
+    public ScriptTemplate asScript() {
+        ParamsBuilder params = paramsBuilder();
+
+        ScriptTemplate script = super.asScript();
+        String template = formatTemplate("{sql}." + StringUtils.underscoreToLowerCamelCase(nameExtractor.name())
+            + "(" + script.template() + ", {})");
+        params.script(script.params())
+              .variable(zoneId().getId());
+
+        return new ScriptTemplate(template, params.build(), dataType());
     }
 
     @Override

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

@@ -6,18 +6,16 @@
 package org.elasticsearch.xpack.sql.expression.function.scalar.datetime;
 
 import org.elasticsearch.xpack.sql.expression.Expression;
-import org.elasticsearch.xpack.sql.expression.FieldAttribute;
 import org.elasticsearch.xpack.sql.expression.function.scalar.datetime.NonIsoDateTimeProcessor.NonIsoDateTimeExtractor;
 import org.elasticsearch.xpack.sql.expression.gen.processor.Processor;
+import org.elasticsearch.xpack.sql.expression.gen.script.ParamsBuilder;
 import org.elasticsearch.xpack.sql.expression.gen.script.ScriptTemplate;
 import org.elasticsearch.xpack.sql.tree.Source;
 import org.elasticsearch.xpack.sql.type.DataType;
 import org.elasticsearch.xpack.sql.util.StringUtils;
 
 import java.time.ZoneId;
-import java.util.Locale;
 
-import static java.lang.String.format;
 import static org.elasticsearch.xpack.sql.expression.gen.script.ParamsBuilder.paramsBuilder;
 
 /*
@@ -33,14 +31,16 @@ abstract class NonIsoDateTimeFunction extends BaseDateTimeFunction {
     }
 
     @Override
-    public ScriptTemplate scriptWithField(FieldAttribute field) {
-        return new ScriptTemplate(
-                formatTemplate(format(Locale.ROOT, "{sql}.%s(doc[{}].value, {})",
-                        StringUtils.underscoreToLowerCamelCase(extractor.name()))),
-                paramsBuilder()
-                  .variable(field.name())
-                  .variable(zoneId().getId()).build(),
-                dataType());
+    public ScriptTemplate asScript() {
+        ParamsBuilder params = paramsBuilder();
+
+        ScriptTemplate script = super.asScript();
+        String template = formatTemplate("{sql}." + StringUtils.underscoreToLowerCamelCase(extractor.name())
+            + "(" + script.template() + ", {})");
+        params.script(script.params())
+              .variable(zoneId().getId());
+
+        return new ScriptTemplate(template, params.build(), dataType());
     }
 
     @Override

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

@@ -7,8 +7,8 @@
 package org.elasticsearch.xpack.sql.expression.function.scalar.datetime;
 
 import org.elasticsearch.xpack.sql.expression.Expression;
-import org.elasticsearch.xpack.sql.expression.FieldAttribute;
 import org.elasticsearch.xpack.sql.expression.gen.processor.Processor;
+import org.elasticsearch.xpack.sql.expression.gen.script.ParamsBuilder;
 import org.elasticsearch.xpack.sql.expression.gen.script.ScriptTemplate;
 import org.elasticsearch.xpack.sql.tree.NodeInfo.NodeCtor2;
 import org.elasticsearch.xpack.sql.tree.Source;
@@ -23,15 +23,17 @@ public class Quarter extends BaseDateTimeFunction {
     public Quarter(Source source, Expression field, ZoneId zoneId) {
         super(source, field, zoneId);
     }
-
+    
     @Override
-    public ScriptTemplate scriptWithField(FieldAttribute field) {
-        return new ScriptTemplate(formatTemplate("{sql}.quarter(doc[{}].value, {})"),
-                paramsBuilder()
-                  .variable(field.name())
-                  .variable(zoneId().getId())
-                  .build(),
-                dataType());
+    public ScriptTemplate asScript() {
+        ParamsBuilder params = paramsBuilder();
+
+        ScriptTemplate script = super.asScript();
+        String template = formatTemplate("{sql}.quarter(" + script.template() + ", {})");
+        params.script(script.params())
+              .variable(zoneId().getId());
+
+        return new ScriptTemplate(template, params.build(), dataType());
     }
 
     @Override

x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/planner/QueryTranslatorTests.java

@@ -24,6 +24,7 @@
 import org.elasticsearch.xpack.sql.expression.function.FunctionRegistry;
 import org.elasticsearch.xpack.sql.expression.function.grouping.Histogram;
 import org.elasticsearch.xpack.sql.expression.function.scalar.Cast;
+import org.elasticsearch.xpack.sql.expression.function.scalar.datetime.DateTimeProcessor.DateTimeExtractor;
 import org.elasticsearch.xpack.sql.expression.function.scalar.math.MathProcessor.MathOperation;
 import org.elasticsearch.xpack.sql.expression.function.scalar.math.Round;
 import org.elasticsearch.xpack.sql.expression.gen.script.ScriptTemplate;
@@ -1160,4 +1161,38 @@ public void testZonedDateTimeInScripts() {
                 + "\"lang\":\"painless\","
                 + "\"params\":{\"v0\":\"date\",\"v1\":\"P1Y\",\"v2\":\"INTERVAL_YEAR\",\"v3\":\"2019-03-11T12:34:56.000Z\"}},"));
     }
+
+    public void testChronoFieldBasedDateTimeFunctionsWithMathIntervalAndGroupBy() {
+        DateTimeExtractor randomFunction = randomValueOtherThan(DateTimeExtractor.YEAR, () -> randomFrom(DateTimeExtractor.values()));
+        PhysicalPlan p = optimizeAndPlan(
+                "SELECT "
+                + randomFunction.name()
+                + "(date + INTERVAL 1 YEAR) FROM test GROUP BY " + randomFunction.name() + "(date + INTERVAL 1 YEAR)");
+        assertEquals(EsQueryExec.class, p.getClass());
+        EsQueryExec eqe = (EsQueryExec) p;
+        assertThat(eqe.queryContainer().toString().replaceAll("\\s+", ""), containsString(
+                "{\"terms\":{\"script\":{\"source\":\"InternalSqlScriptUtils.dateTimeChrono("
+                + "InternalSqlScriptUtils.add(InternalSqlScriptUtils.docValue(doc,params.v0),"
+                + "InternalSqlScriptUtils.intervalYearMonth(params.v1,params.v2)),params.v3,params.v4)\","
+                + "\"lang\":\"painless\",\"params\":{\"v0\":\"date\",\"v1\":\"P1Y\",\"v2\":\"INTERVAL_YEAR\","
+                + "\"v3\":\"Z\",\"v4\":\"" + randomFunction.chronoField().name() + "\"}},\"missing_bucket\":true,"
+                + "\"value_type\":\"long\",\"order\":\"asc\"}}}]}}}}"));
+    }
+
+    public void testDateTimeFunctionsWithMathIntervalAndGroupBy() {
+        String[] functions = new String[] {"DAY_NAME", "MONTH_NAME", "DAY_OF_WEEK", "WEEK_OF_YEAR", "QUARTER"};
+        String[] scriptMethods = new String[] {"dayName", "monthName", "dayOfWeek", "weekOfYear", "quarter"};
+        int pos = randomIntBetween(0, functions.length - 1);
+        PhysicalPlan p = optimizeAndPlan(
+                "SELECT "
+                + functions[pos]
+                + "(date + INTERVAL 1 YEAR) FROM test GROUP BY " + functions[pos] + "(date + INTERVAL 1 YEAR)");
+        assertEquals(EsQueryExec.class, p.getClass());
+        EsQueryExec eqe = (EsQueryExec) p;
+        assertThat(eqe.queryContainer().toString().replaceAll("\\s+", ""), containsString(
+                "{\"terms\":{\"script\":{\"source\":\"InternalSqlScriptUtils." + scriptMethods[pos] 
+                + "(InternalSqlScriptUtils.add(InternalSqlScriptUtils.docValue(doc,params.v0),"
+                + "InternalSqlScriptUtils.intervalYearMonth(params.v1,params.v2)),params.v3)\",\"lang\":\"painless\","
+                + "\"params\":{\"v0\":\"date\",\"v1\":\"P1Y\",\"v2\":\"INTERVAL_YEAR\",\"v3\":\"Z\"}},\"missing_bucket\":true,"));
+    }
 }