SQL: Added support for string manipulating functions with more than one parameter

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

@@ -68,6 +68,15 @@
 import org.elasticsearch.xpack.sql.expression.function.scalar.string.RTrim;
 import org.elasticsearch.xpack.sql.expression.function.scalar.string.Space;
 import org.elasticsearch.xpack.sql.expression.function.scalar.string.UCase;
+import org.elasticsearch.xpack.sql.expression.function.scalar.string.Concat;
+import org.elasticsearch.xpack.sql.expression.function.scalar.string.Insert;
+import org.elasticsearch.xpack.sql.expression.function.scalar.string.Left;
+import org.elasticsearch.xpack.sql.expression.function.scalar.string.Locate;
+import org.elasticsearch.xpack.sql.expression.function.scalar.string.Position;
+import org.elasticsearch.xpack.sql.expression.function.scalar.string.Repeat;
+import org.elasticsearch.xpack.sql.expression.function.scalar.string.Replace;
+import org.elasticsearch.xpack.sql.expression.function.scalar.string.Right;
+import org.elasticsearch.xpack.sql.expression.function.scalar.string.Substring;
 import org.elasticsearch.xpack.sql.parser.ParsingException;
 import org.elasticsearch.xpack.sql.tree.Location;
 import org.elasticsearch.xpack.sql.util.StringUtils;
@@ -154,6 +163,15 @@ public class FunctionRegistry {
             def(LTrim.class, LTrim::new),
             def(RTrim.class, RTrim::new),
             def(Space.class, Space::new),
+            def(Concat.class, Concat::new),
+            def(Insert.class, Insert::new),
+            def(Left.class, Left::new),
+            def(Locate.class, Locate::new),
+            def(Position.class, Position::new),
+            def(Repeat.class, Repeat::new),
+            def(Replace.class, Replace::new),
+            def(Right.class, Right::new),
+            def(Substring.class, Substring::new),
             def(UCase.class, UCase::new),
         // Special
             def(Score.class, Score::new)));
@@ -337,6 +355,47 @@ private static FunctionDefinition def(Class<? extends Function> function, Functi
     private interface FunctionBuilder {
         Function build(Location location, List<Expression> children, boolean distinct, TimeZone tz);
     }
+
+    @SuppressWarnings("overloads")  // These are ambiguous if you aren't using ctor references but we always do
+    static <T extends Function> FunctionDefinition def(Class<T> function,
+            ThreeParametersFunctionBuilder<T> ctorRef, String... aliases) {
+        FunctionBuilder builder = (location, children, distinct, tz) -> {
+            boolean isLocateFunction = function.isAssignableFrom(Locate.class);
+            if (isLocateFunction && (children.size() > 3 || children.size() < 2)) {
+                throw new IllegalArgumentException("expects two or three arguments");
+            } else if (!isLocateFunction && children.size() != 3) {
+                throw new IllegalArgumentException("expects exactly three arguments");
+            }
+            if (distinct) {
+                throw new IllegalArgumentException("does not support DISTINCT yet it was specified");
+            }
+            return ctorRef.build(location, children.get(0), children.get(1), children.size() == 3 ? children.get(2) : null);
+        };
+        return def(function, builder, false, aliases);
+    }
+
+    interface ThreeParametersFunctionBuilder<T> {
+        T build(Location location, Expression source, Expression exp1, Expression exp2);
+    }
+
+    @SuppressWarnings("overloads")  // These are ambiguous if you aren't using ctor references but we always do
+    static <T extends Function> FunctionDefinition def(Class<T> function,
+            FourParametersFunctionBuilder<T> ctorRef, String... aliases) {
+        FunctionBuilder builder = (location, children, distinct, tz) -> {
+            if (children.size() != 4) {
+                throw new IllegalArgumentException("expects exactly four arguments");
+            }
+            if (distinct) {
+                throw new IllegalArgumentException("does not support DISTINCT yet it was specified");
+            }
+            return ctorRef.build(location, children.get(0), children.get(1), children.get(2), children.get(3));
+        };
+        return def(function, builder, false, aliases);
+    }
+
+    interface FourParametersFunctionBuilder<T> {
+        T build(Location location, Expression source, Expression exp1, Expression exp2, Expression exp3);
+    }
 
     private static String normalize(String name) {
         // translate CamelCase to camel_case

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

@@ -1,5 +1,5 @@
 /*
- * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+* 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.
  */
@@ -18,6 +18,13 @@
 import org.elasticsearch.xpack.sql.expression.function.scalar.processor.runtime.HitExtractorProcessor;
 import org.elasticsearch.xpack.sql.expression.function.scalar.processor.runtime.Processor;
 import org.elasticsearch.xpack.sql.expression.function.scalar.string.StringProcessor;
+import org.elasticsearch.xpack.sql.expression.function.scalar.string.BinaryStringNumericProcessor;
+import org.elasticsearch.xpack.sql.expression.function.scalar.string.BinaryStringStringProcessor;
+import org.elasticsearch.xpack.sql.expression.function.scalar.string.ConcatFunctionProcessor;
+import org.elasticsearch.xpack.sql.expression.function.scalar.string.InsertFunctionProcessor;
+import org.elasticsearch.xpack.sql.expression.function.scalar.string.LocateFunctionProcessor;
+import org.elasticsearch.xpack.sql.expression.function.scalar.string.ReplaceFunctionProcessor;
+import org.elasticsearch.xpack.sql.expression.function.scalar.string.SubstringFunctionProcessor;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -49,6 +56,13 @@ public static List<NamedWriteableRegistry.Entry> getNamedWriteables() {
         entries.add(new Entry(Processor.class, MathProcessor.NAME, MathProcessor::new));
         // string
         entries.add(new Entry(Processor.class, StringProcessor.NAME, StringProcessor::new));
+        entries.add(new Entry(Processor.class, BinaryStringNumericProcessor.NAME, BinaryStringNumericProcessor::new));
+        entries.add(new Entry(Processor.class, BinaryStringStringProcessor.NAME, BinaryStringStringProcessor::new));
+        entries.add(new Entry(Processor.class, ConcatFunctionProcessor.NAME, ConcatFunctionProcessor::new));
+        entries.add(new Entry(Processor.class, InsertFunctionProcessor.NAME, InsertFunctionProcessor::new));
+        entries.add(new Entry(Processor.class, LocateFunctionProcessor.NAME, LocateFunctionProcessor::new));
+        entries.add(new Entry(Processor.class, ReplaceFunctionProcessor.NAME, ReplaceFunctionProcessor::new));
+        entries.add(new Entry(Processor.class, SubstringFunctionProcessor.NAME, SubstringFunctionProcessor::new));
         return entries;
     }
 }

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

@@ -0,0 +1,91 @@
+/*
+ * 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.string;
+
+import org.elasticsearch.xpack.sql.expression.Expression;
+import org.elasticsearch.xpack.sql.expression.FieldAttribute;
+import org.elasticsearch.xpack.sql.expression.function.scalar.BinaryScalarFunction;
+import org.elasticsearch.xpack.sql.expression.function.scalar.script.ScriptTemplate;
+import org.elasticsearch.xpack.sql.tree.Location;
+import org.elasticsearch.xpack.sql.type.DataType;
+import org.elasticsearch.xpack.sql.util.StringUtils;
+
+import java.util.Locale;
+import java.util.Objects;
+import java.util.function.BiFunction;
+
+import static java.lang.String.format;
+import static org.elasticsearch.xpack.sql.expression.function.scalar.script.ParamsBuilder.paramsBuilder;
+import static org.elasticsearch.xpack.sql.expression.function.scalar.script.ScriptTemplate.formatTemplate;
+
+/**
+ * Base class for binary functions that have the first parameter a string, the second parameter a number
+ * or a string and the result can be a string or a number.
+ */
+public abstract class BinaryStringFunction<T,R> extends BinaryScalarFunction {
+
+    protected BinaryStringFunction(Location location, Expression left, Expression right) {
+        super(location, left, right);
+    }
+
+    protected abstract BiFunction<String, T, R> operation();
+
+    @Override
+    protected TypeResolution resolveType() {
+        if (!childrenResolved()) {
+            return new TypeResolution("Unresolved children");
+        }
+
+        if (!left().dataType().isString()) {
+            return new TypeResolution("'%s' requires first parameter to be a string type, received %s", functionName(), left().dataType());
+        }
+
+        return resolveSecondParameterInputType(right().dataType());
+    }
+
+    protected abstract TypeResolution resolveSecondParameterInputType(DataType inputType);
+
+    @Override
+    public Object fold() {
+        @SuppressWarnings("unchecked")
+        T fold = (T) right().fold();
+        return operation().apply((String) left().fold(), fold);
+    }
+
+    @Override
+    protected ScriptTemplate asScriptFrom(ScriptTemplate leftScript, ScriptTemplate rightScript) {
+        // basically, transform the script to InternalSqlScriptUtils.[function_name](function_or_field1, function_or_field2)
+        return new ScriptTemplate(format(Locale.ROOT, formatTemplate("{sql}.%s(%s,%s)"), 
+                StringUtils.underscoreToLowerCamelCase(operation().toString()), 
+                leftScript.template(), 
+                rightScript.template()),
+                paramsBuilder()
+                    .script(leftScript.params()).script(rightScript.params())
+                    .build(), dataType());
+    }
+
+    @Override
+    protected ScriptTemplate asScriptFrom(FieldAttribute field) {
+        return new ScriptTemplate(formatScript("doc[{}].value"),
+                paramsBuilder().variable(field.isInexact() ? field.exactAttribute().name() : field.name()).build(),
+                dataType());
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(left(), right());
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+        if (obj == null || obj.getClass() != getClass()) {
+            return false;
+        }
+        BinaryStringFunction<?,?> other = (BinaryStringFunction<?,?>) obj;
+        return Objects.equals(other.left(), left())
+            && Objects.equals(other.right(), right());
+    }
+}

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

@@ -0,0 +1,32 @@
+/*
+ * 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.string;
+
+import org.elasticsearch.xpack.sql.expression.Expression;
+import org.elasticsearch.xpack.sql.tree.Location;
+import org.elasticsearch.xpack.sql.type.DataType;
+
+/**
+ * A binary string function with a numeric second parameter and a string result
+ */
+public abstract class BinaryStringNumericFunction extends BinaryStringFunction<Number, String> {
+
+    public BinaryStringNumericFunction(Location location, Expression left, Expression right) {
+        super(location, left, right);
+    }
+
+    @Override
+    protected TypeResolution resolveSecondParameterInputType(DataType inputType) {
+        return inputType.isNumeric() ? 
+                TypeResolution.TYPE_RESOLVED : 
+                new TypeResolution("'%s' requires second parameter to be a numeric type, received %s", functionName(), inputType);
+    }
+
+    @Override
+    public DataType dataType() {
+        return DataType.KEYWORD;
+    }
+}

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

@@ -0,0 +1,88 @@
+/*
+ * 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.string;
+
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.xpack.sql.SqlIllegalArgumentException;
+import org.elasticsearch.xpack.sql.expression.function.scalar.processor.runtime.Processor;
+import org.elasticsearch.xpack.sql.expression.function.scalar.string.BinaryStringNumericProcessor.BinaryStringNumericOperation;
+
+import java.io.IOException;
+import java.util.function.BiFunction;
+
+public class BinaryStringNumericProcessor extends BinaryStringProcessor<BinaryStringNumericOperation, Number, String> {
+
+    public static final String NAME = "sn";
+
+    public BinaryStringNumericProcessor(StreamInput in) throws IOException {
+        super(in, i -> i.readEnum(BinaryStringNumericOperation.class));
+    }
+
+    public BinaryStringNumericProcessor(Processor left, Processor right, BinaryStringNumericOperation operation) {
+        super(left, right, operation);
+    }
+
+    public enum BinaryStringNumericOperation implements BiFunction<String, Number, String> {
+        LEFT((s,c) -> {
+            int i = c.intValue();
+            if (i < 0) return "";
+            return i > s.length() ? s : s.substring(0, i);
+        }),
+        RIGHT((s,c) -> {
+            int i = c.intValue();
+            if (i < 0) return "";
+            return i > s.length() ? s : s.substring(s.length() - i);
+        }),
+        REPEAT((s,c) -> {
+            int i = c.intValue();
+            if (i <= 0) return null;
+
+            StringBuilder sb = new StringBuilder(s.length() * i);
+            for (int j = 0; j < i; j++) {
+                sb.append(s);
+            }
+            return sb.toString();
+        });
+
+        BinaryStringNumericOperation(BiFunction<String, Number, String> op) {
+            this.op = op;
+        }
+
+        private final BiFunction<String, Number, String> op;
+
+        @Override
+        public String apply(String stringExp, Number count) {
+            return op.apply(stringExp, count);
+        }
+    }
+
+    @Override
+    protected void doWrite(StreamOutput out) throws IOException {
+        out.writeEnum(operation());
+    }
+
+    @Override
+    protected Object doProcess(Object left, Object right) {
+        if (left == null || right == null) {
+            return null;
+        }
+        if (!(left instanceof String || left instanceof Character)) {
+            throw new SqlIllegalArgumentException("A string/char is required; received [{}]", left);
+        }
+        if (!(right instanceof Number)) {
+            throw new SqlIllegalArgumentException("A number is required; received [{}]", right);
+        }
+
+        return operation().apply(left instanceof Character ? left.toString() : (String) left, (Number) right);
+    }
+
+    @Override
+    public String getWriteableName() {
+        return NAME;
+    }
+
+}