SQL: Refactor binary date time functions (#47786)

Refactor DateTrunc and DatePart to use separate Pipe classes which
allows the removal of the BinaryDateOperation enum.

(cherry picked from commit a6075e7)

Author: matriv

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

@@ -20,19 +20,15 @@
 import static org.elasticsearch.common.logging.LoggerMessageFormat.format;
 import static org.elasticsearch.xpack.sql.expression.TypeResolutions.isDate;
 import static org.elasticsearch.xpack.sql.expression.TypeResolutions.isString;
-import static org.elasticsearch.xpack.sql.expression.function.scalar.datetime.BinaryDateTimeProcessor.BinaryDateOperation;
 import static org.elasticsearch.xpack.sql.expression.gen.script.ParamsBuilder.paramsBuilder;
 
 public abstract class BinaryDateTimeFunction extends BinaryScalarFunction {
 
     private final ZoneId zoneId;
-    private final BinaryDateOperation operation;
 
-    public BinaryDateTimeFunction(Source source, Expression datePart, Expression timestamp, ZoneId zoneId,
-                                  BinaryDateOperation operation) {
+    public BinaryDateTimeFunction(Source source, Expression datePart, Expression timestamp, ZoneId zoneId) {
         super(source, datePart, timestamp);
         this.zoneId = zoneId;
-        this.operation = operation;
     }
 
     @Override
@@ -47,7 +43,7 @@ protected TypeResolution resolveType() {
             if (datePartValue != null && resolveDateTimeField(datePartValue) == false) {
                 List<String> similar = findSimilarDateTimeFields(datePartValue);
                 if (similar.isEmpty()) {
-                    return new TypeResolution(format(null, "first argument of [{}] must be one of {} or their aliases, found value [{}]",
+                    return new TypeResolution(format(null, "first argument of [{}] must be one of {} or their aliases; found value [{}]",
                         sourceText(),
                         validDateTimeFieldValues(),
                         Expressions.name(left())));
@@ -78,9 +74,11 @@ public ZoneId zoneId() {
 
     @Override
     protected Pipe makePipe() {
-        return new BinaryDateTimePipe(source(), this, Expressions.pipe(left()), Expressions.pipe(right()), zoneId, operation);
+        return createPipe(Expressions.pipe(left()), Expressions.pipe(right()), zoneId);
     }
 
+    protected abstract Pipe createPipe(Pipe left, Pipe right, ZoneId zoneId);
+
     @Override
     public Nullability nullable() {
         return Nullability.TRUE;
@@ -101,7 +99,7 @@ protected ScriptTemplate asScriptFrom(ScriptTemplate leftScript, ScriptTemplate
 
     @Override
     public int hashCode() {
-        return Objects.hash(super.hashCode(), zoneId, operation);
+        return Objects.hash(super.hashCode(), zoneId);
     }
 
     @Override
@@ -116,6 +114,6 @@ public boolean equals(Object o) {
             return false;
         }
         BinaryDateTimeFunction that = (BinaryDateTimeFunction) o;
-        return zoneId.equals(that.zoneId) && operation == that.operation;
+        return zoneId.equals(that.zoneId);
     }
 }

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

@@ -9,50 +9,34 @@
 import org.elasticsearch.xpack.sql.expression.gen.pipeline.BinaryPipe;
 import org.elasticsearch.xpack.sql.expression.gen.pipeline.Pipe;
 import org.elasticsearch.xpack.sql.expression.gen.processor.Processor;
-import org.elasticsearch.xpack.sql.tree.NodeInfo;
 import org.elasticsearch.xpack.sql.tree.Source;
 
 import java.time.ZoneId;
 import java.util.Objects;
 
-public class BinaryDateTimePipe extends BinaryPipe {
+public abstract class BinaryDateTimePipe extends BinaryPipe {
 
     private final ZoneId zoneId;
-    private final BinaryDateTimeProcessor.BinaryDateOperation operation;
 
-    public BinaryDateTimePipe(Source source, Expression expression, Pipe left, Pipe right, ZoneId zoneId,
-                              BinaryDateTimeProcessor.BinaryDateOperation operation) {
+    public BinaryDateTimePipe(Source source, Expression expression, Pipe left, Pipe right, ZoneId zoneId) {
         super(source, expression, left, right);
         this.zoneId = zoneId;
-        this.operation = operation;
     }
 
     ZoneId zoneId() {
         return zoneId;
     }
 
-    BinaryDateTimeProcessor.BinaryDateOperation operation() {
-        return operation;
-    }
-
-    @Override
-    protected NodeInfo<BinaryDateTimePipe> info() {
-        return NodeInfo.create(this, BinaryDateTimePipe::new, expression(), left(), right(), zoneId, operation);
-    }
-
-    @Override
-    protected BinaryPipe replaceChildren(Pipe left, Pipe right) {
-        return new BinaryDateTimePipe(source(), expression(), left, right, zoneId, operation);
-    }
-
     @Override
     public Processor asProcessor() {
-        return BinaryDateTimeProcessor.asProcessor(operation, left().asProcessor(), right().asProcessor(), zoneId);
+        return makeProcessor(left().asProcessor(), right().asProcessor(), zoneId);
     }
 
+    protected abstract Processor makeProcessor(Processor left, Processor right, ZoneId zoneId);
+
     @Override
     public int hashCode() {
-        return Objects.hash(super.hashCode(), zoneId, operation);
+        return Objects.hash(super.hashCode(), zoneId);
     }
 
     @Override
@@ -67,7 +51,6 @@ public boolean equals(Object o) {
             return false;
         }
         BinaryDateTimePipe that = (BinaryDateTimePipe) o;
-        return Objects.equals(zoneId, that.zoneId) &&
-            operation == that.operation;
+        return Objects.equals(zoneId, that.zoneId);
     }
 }

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

@@ -15,16 +15,8 @@
 import java.time.ZoneId;
 import java.util.Objects;
 
-import static org.elasticsearch.xpack.sql.expression.function.scalar.datetime.BinaryDateTimeProcessor.BinaryDateOperation.TRUNC;
-
 public abstract class BinaryDateTimeProcessor extends BinaryProcessor {
 
-    // TODO: Remove and in favour of inheritance (subclasses which implement abstract methods)
-    public enum BinaryDateOperation {
-        TRUNC,
-        PART;
-    }
-
     private final ZoneId zoneId;
 
     public BinaryDateTimeProcessor(Processor source1, Processor source2, ZoneId zoneId) {
@@ -48,28 +40,24 @@ ZoneId zoneId() {
     @Override
     protected abstract Object doProcess(Object left, Object right);
 
-    public static BinaryDateTimeProcessor asProcessor(BinaryDateOperation operation, Processor left, Processor right, ZoneId zoneId) {
-        if (operation == TRUNC) {
-            return new DateTruncProcessor(left, right, zoneId);
-        } else {
-            return new DatePartProcessor(left, right, zoneId);
-        }
-    }
-
     @Override
     public int hashCode() {
-        return Objects.hash(zoneId);
+        return Objects.hash(left(), right(), zoneId);
     }
 
     @Override
-    public boolean equals(Object o) {
-        if (this == o) {
+    public boolean equals(Object obj) {
+        if (this == obj) {
             return true;
         }
-        if (o == null || getClass() != o.getClass()) {
+
+        if (obj == null || getClass() != obj.getClass()) {
             return false;
         }
-        BinaryDateTimeProcessor that = (BinaryDateTimeProcessor) o;
-        return zoneId.equals(that.zoneId);
+
+        BinaryDateTimeProcessor other = (BinaryDateTimeProcessor) obj;
+        return Objects.equals(left(), other.left())
+            && Objects.equals(right(), other.right())
+            && Objects.equals(zoneId(), other.zoneId());
     }
 }

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

@@ -9,6 +9,7 @@
 import org.elasticsearch.xpack.sql.expression.Nullability;
 import org.elasticsearch.xpack.sql.expression.function.scalar.BinaryScalarFunction;
 import org.elasticsearch.xpack.sql.expression.function.scalar.datetime.DateTimeProcessor.DateTimeExtractor;
+import org.elasticsearch.xpack.sql.expression.gen.pipeline.Pipe;
 import org.elasticsearch.xpack.sql.tree.NodeInfo;
 import org.elasticsearch.xpack.sql.tree.Source;
 import org.elasticsearch.xpack.sql.type.DataType;
@@ -78,7 +79,7 @@ public Integer extract(ZonedDateTime dateTime) {
     }
 
     public DatePart(Source source, Expression truncateTo, Expression timestamp, ZoneId zoneId) {
-        super(source, truncateTo, timestamp, zoneId, BinaryDateTimeProcessor.BinaryDateOperation.PART);
+        super(source, truncateTo, timestamp, zoneId);
     }
 
     @Override
@@ -102,23 +103,28 @@ public Nullability nullable() {
     }
 
     @Override
-    protected boolean resolveDateTimeField(String dateTimeField) {
-        return Part.resolve(dateTimeField) != null;
+    protected String scriptMethodName() {
+        return "datePart";
     }
 
     @Override
-    protected List<String> findSimilarDateTimeFields(String dateTimeField) {
-        return Part.findSimilar(dateTimeField);
+    public Object fold() {
+        return DatePartProcessor.process(left().fold(), right().fold(), zoneId());
     }
 
     @Override
-    protected String scriptMethodName() {
-        return "datePart";
+    protected Pipe createPipe(Pipe left, Pipe right, ZoneId zoneId) {
+        return new DatePartPipe(source(), this, left, right, zoneId);
     }
 
     @Override
-    public Object fold() {
-        return DatePartProcessor.process(left().fold(), right().fold(), zoneId());
+    protected boolean resolveDateTimeField(String dateTimeField) {
+        return Part.resolve(dateTimeField) != null;
+    }
+
+    @Override
+    protected List<String> findSimilarDateTimeFields(String dateTimeField) {
+        return Part.findSimilar(dateTimeField);
     }
 
     @Override

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

@@ -0,0 +1,36 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+package org.elasticsearch.xpack.sql.expression.function.scalar.datetime;
+
+import org.elasticsearch.xpack.sql.expression.Expression;
+import org.elasticsearch.xpack.sql.expression.gen.pipeline.Pipe;
+import org.elasticsearch.xpack.sql.expression.gen.processor.Processor;
+import org.elasticsearch.xpack.sql.tree.NodeInfo;
+import org.elasticsearch.xpack.sql.tree.Source;
+
+import java.time.ZoneId;
+
+public class DatePartPipe extends BinaryDateTimePipe {
+
+    public DatePartPipe(Source source, Expression expression, Pipe left, Pipe right, ZoneId zoneId) {
+        super(source, expression, left, right, zoneId);
+    }
+
+    @Override
+    protected NodeInfo<DatePartPipe> info() {
+        return NodeInfo.create(this, DatePartPipe::new, expression(), left(), right(), zoneId());
+    }
+
+    @Override
+    protected DatePartPipe replaceChildren(Pipe left, Pipe right) {
+        return new DatePartPipe(source(), expression(), left, right, zoneId());
+    }
+
+    @Override
+    protected Processor makeProcessor(Processor left, Processor right, ZoneId zoneId) {
+        return new DatePartProcessor(left, right, zoneId);
+    }
+}

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

@@ -34,36 +34,36 @@ public String getWriteableName() {
     }
 
     @Override
-    protected Object doProcess(Object left, Object right) {
-        return process(left, right, zoneId());
+    protected Object doProcess(Object part, Object timestamp) {
+        return process(part, timestamp, zoneId());
     }
 
     /**
      * Used in Painless scripting
      */
-    public static Object process(Object source1, Object source2, ZoneId zoneId) {
-        if (source1 == null || source2 == null) {
+    public static Object process(Object part, Object timestamp, ZoneId zoneId) {
+        if (part == null || timestamp == null) {
             return null;
         }
-        if (source1 instanceof String == false) {
-            throw new SqlIllegalArgumentException("A string is required; received [{}]", source1);
+        if (part instanceof String == false) {
+            throw new SqlIllegalArgumentException("A string is required; received [{}]", part);
         }
-        Part datePartField = Part.resolve((String) source1);
+        Part datePartField = Part.resolve((String) part);
         if (datePartField == null) {
-            List<String> similar = Part.findSimilar((String) source1);
+            List<String> similar = Part.findSimilar((String) part);
             if (similar.isEmpty()) {
                 throw new SqlIllegalArgumentException("A value of {} or their aliases is required; received [{}]",
-                    Part.values(), source1);
+                    Part.values(), part);
             } else {
                 throw new SqlIllegalArgumentException("Received value [{}] is not valid date part for extraction; " +
-                    "did you mean {}?", source1, similar);
+                    "did you mean {}?", part, similar);
             }
         }
 
-        if (source2 instanceof ZonedDateTime == false) {
-            throw new SqlIllegalArgumentException("A date/datetime is required; received [{}]", source2);
+        if (timestamp instanceof ZonedDateTime == false) {
+            throw new SqlIllegalArgumentException("A date/datetime is required; received [{}]", timestamp);
         }
 
-        return datePartField.extract(((ZonedDateTime) source2).withZoneSameInstant(zoneId));
+        return datePartField.extract(((ZonedDateTime) timestamp).withZoneSameInstant(zoneId));
     }
 }