Merge from main

Signed-off-by: Andy Kwok <[email protected]>

Author: andy-k-improving

core/src/main/java/org/opensearch/sql/data/type/ExprType.java

@@ -9,6 +9,7 @@
 
 import java.util.Arrays;
 import java.util.List;
+import java.util.Optional;
 import org.opensearch.sql.data.model.ExprValue;
 import org.opensearch.sql.expression.Expression;
 
@@ -54,4 +55,17 @@ default List<ExprType> getParent() {
   default String legacyTypeName() {
     return typeName();
   }
+
+  /** Get the original path. Types like alias type will set the actual path in field property. */
+  default Optional<String> getOriginalPath() {
+    return Optional.empty();
+  }
+
+  /**
+   * Get the original path. Types like alias type should be derived from the type of the original
+   * field.
+   */
+  default ExprType getOriginalExprType() {
+    return this;
+  }
 }

core/src/main/java/org/opensearch/sql/expression/DSL.java

@@ -985,6 +985,10 @@ public static FunctionExpression utc_timestamp(
     return compile(functionProperties, BuiltinFunctionName.UTC_TIMESTAMP, args);
   }
 
+  public static FunctionExpression geoip(Expression... args) {
+    return compile(FunctionProperties.None, BuiltinFunctionName.GEOIP, args);
+  }
+
   public static FunctionExpression json_function(Expression value) {
     return compile(FunctionProperties.None, BuiltinFunctionName.JSON, value);
   }

core/src/main/java/org/opensearch/sql/expression/ReferenceExpression.java

@@ -23,6 +23,8 @@
 public class ReferenceExpression implements Expression {
   @Getter private final String attr;
 
+  @Getter private final String rawPath;
+
   @Getter private final List<String> paths;
 
   private final ExprType type;
@@ -36,8 +38,11 @@ public class ReferenceExpression implements Expression {
   public ReferenceExpression(String ref, ExprType type) {
     this.attr = ref;
     // Todo. the define of paths need to be redefined after adding multiple index/variable support.
-    this.paths = Arrays.asList(ref.split("\\."));
-    this.type = type;
+    // For AliasType, the actual path is set in the property of `path` and the type is derived
+    // from the type of field on that path; Otherwise, use ref itself as the path
+    this.rawPath = type.getOriginalPath().orElse(ref);
+    this.paths = Arrays.asList(rawPath.split("\\."));
+    this.type = type.getOriginalExprType();
   }
 
   @Override

core/src/main/java/org/opensearch/sql/expression/datetime/DateTimeFunctions.java

@@ -63,6 +63,7 @@
 import java.time.format.DateTimeParseException;
 import java.time.format.TextStyle;
 import java.time.temporal.ChronoUnit;
+import java.time.temporal.IsoFields;
 import java.time.temporal.TemporalAmount;
 import java.util.Locale;
 import java.util.Map;
@@ -123,7 +124,8 @@ public class DateTimeFunctions {
           .put("MINUTE", "mm")
           .put("HOUR", "HH")
           .put("DAY", "dd")
-          .put("WEEK", "w")
+          // removing "WEEK" to standardize the extract
+          // .put("WEEK", "w")
           .put("MONTH", "MM")
           .put("YEAR", "yyyy")
           .put("SECOND_MICROSECOND", "ssSSSSSS")
@@ -1436,6 +1438,12 @@ private ExprValue exprDayOfYear(ExprValue date) {
   public ExprLongValue formatExtractFunction(ExprValue part, ExprValue timestamp) {
     String partName = part.stringValue().toUpperCase();
     LocalDateTime arg = timestamp.timestampValue().atZone(ZoneOffset.UTC).toLocalDateTime();
+
+    // Override "Week" to use the IsoFields week-of-week-based-year format
+    if (partName.equals("WEEK")) {
+      return new ExprLongValue(arg.get(IsoFields.WEEK_OF_WEEK_BASED_YEAR));
+    }
+
     String text =
         arg.format(DateTimeFormatter.ofPattern(extract_formats.get(partName), Locale.ENGLISH));
 

core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java

@@ -210,6 +210,9 @@ public enum BuiltinFunctionName {
   JSON_EXTRACT(FunctionName.of("json_extract")),
   JSON_SET(FunctionName.of("json_set")),
 
+  /** GEOSPATIAL Functions. */
+  GEOIP(FunctionName.of("geoip")),
+
   /** NULL Test. */
   IS_NULL(FunctionName.of("is null")),
   IS_NOT_NULL(FunctionName.of("is not null")),

core/src/main/java/org/opensearch/sql/expression/function/OpenSearchFunctions.java

@@ -173,4 +173,54 @@ public String toString() {
       return String.format("%s(%s)", functionName, String.join(", ", args));
     }
   }
+
+  /**
+   * Static class to identify functional Expression which specifically designed for OpenSearch
+   * storage runtime
+   */
+  public static class OpenSearchExecutableFunction extends FunctionExpression {
+    private final FunctionName functionName;
+    private final List<Expression> arguments;
+    private final ExprType returnType;
+
+    public OpenSearchExecutableFunction(
+        FunctionName functionName, List<Expression> arguments, ExprType returnType) {
+      super(functionName, arguments);
+      this.functionName = functionName;
+      this.arguments = arguments;
+      this.returnType = returnType;
+    }
+
+    @Override
+    public ExprValue valueOf(Environment<Expression, ExprValue> valueEnv) {
+      throw new UnsupportedOperationException(
+          String.format(
+              "OpenSearch defined function [%s] is only supported in Eval operation.",
+              functionName));
+    }
+
+    @Override
+    public ExprType type() {
+      return returnType;
+    }
+
+    /**
+     * Util method to generate probe implementation with given list of argument types, with marker
+     * class `OpenSearchFunction` to annotate this is an OpenSearch specific expression.
+     *
+     * @param returnType return type.
+     * @return Binary Function Implementation.
+     */
+    public static SerializableFunction<FunctionName, Pair<FunctionSignature, FunctionBuilder>>
+        openSearchImpl(ExprType returnType, List<ExprType> args) {
+      return functionName -> {
+        FunctionSignature functionSignature = new FunctionSignature(functionName, args);
+        FunctionBuilder functionBuilder =
+            (functionProperties, arguments) ->
+                new OpenSearchFunctions.OpenSearchExecutableFunction(
+                    functionName, arguments, returnType);
+        return Pair.of(functionSignature, functionBuilder);
+      };
+    }
+  }
 }

core/src/main/java/org/opensearch/sql/expression/ip/IPFunctions.java

@@ -11,8 +11,10 @@
 import static org.opensearch.sql.expression.function.FunctionDSL.define;
 import static org.opensearch.sql.expression.function.FunctionDSL.impl;
 import static org.opensearch.sql.expression.function.FunctionDSL.nullMissingHandling;
+import static org.opensearch.sql.expression.function.OpenSearchFunctions.OpenSearchExecutableFunction.openSearchImpl;
 
 import inet.ipaddr.IPAddress;
+import java.util.Arrays;
 import lombok.experimental.UtilityClass;
 import org.opensearch.sql.data.model.ExprValue;
 import org.opensearch.sql.data.model.ExprValueUtils;
@@ -28,6 +30,7 @@ public class IPFunctions {
 
   public void register(BuiltinFunctionRepository repository) {
     repository.register(cidrmatch());
+    repository.register(geoIp());
   }
 
   private DefaultFunctionResolver cidrmatch() {
@@ -57,4 +60,16 @@ private ExprValue exprCidrMatch(ExprValue addressExprValue, ExprValue rangeExprV
         ? ExprValueUtils.LITERAL_FALSE
         : ExprValueUtils.LITERAL_TRUE;
   }
+
+  /**
+   * To register all method signatures related to geoip( ) expression under eval.
+   *
+   * @return Resolver for geoip( ) expression.
+   */
+  private DefaultFunctionResolver geoIp() {
+    return define(
+        BuiltinFunctionName.GEOIP.getName(),
+        openSearchImpl(BOOLEAN, Arrays.asList(STRING, STRING)),
+        openSearchImpl(BOOLEAN, Arrays.asList(STRING, STRING, STRING)));
+  }
 }

core/src/main/java/org/opensearch/sql/planner/physical/EvalOperator.java

@@ -87,7 +87,7 @@ public ExprValue next() {
    * @param env {@link Environment}
    * @return The mapping of reference and {@link ExprValue} for each expression.
    */
-  private Map<String, ExprValue> eval(Environment<Expression, ExprValue> env) {
+  protected Map<String, ExprValue> eval(Environment<Expression, ExprValue> env) {
     Map<String, ExprValue> evalResultMap = new LinkedHashMap<>();
     for (Pair<ReferenceExpression, Expression> pair : expressionList) {
       ReferenceExpression var = pair.getKey();

core/src/main/java/org/opensearch/sql/utils/JsonUtils.java

@@ -46,9 +46,9 @@ public static ExprValue isValidJson(ExprValue jsonExprValue) {
 
     try {
       objectMapper.readTree(jsonExprValue.stringValue());
-      return LITERAL_TRUE;
+        return ExprValueUtils.LITERAL_TRUE;
     } catch (JsonProcessingException e) {
-      return LITERAL_FALSE;
+        return ExprValueUtils.LITERAL_FALSE;
     }
   }
 

core/src/test/java/org/opensearch/sql/data/type/ExprTypeTest.java

@@ -24,6 +24,7 @@
 import static org.opensearch.sql.data.type.ExprCoreType.UNDEFINED;
 import static org.opensearch.sql.data.type.ExprCoreType.UNKNOWN;
 
+import java.util.Optional;
 import org.hamcrest.Matchers;
 import org.junit.jupiter.api.Test;
 
@@ -85,4 +86,16 @@ void defaultLegacyTypeName() {
     final ExprType exprType = () -> "dummy";
     assertEquals("dummy", exprType.legacyTypeName());
   }
+
+  @Test
+  void getOriginalPath() {
+    final ExprType exprType = () -> "dummy";
+    assertEquals(Optional.empty(), exprType.getOriginalPath());
+  }
+
+  @Test
+  void getOriginalExprType() {
+    final ExprType exprType = () -> "dummy";
+    assertEquals(exprType, exprType.getOriginalExprType());
+  }
 }

core/src/test/java/org/opensearch/sql/expression/datetime/ExtractTest.java

@@ -5,16 +5,21 @@
 
 package org.opensearch.sql.expression.datetime;
 
-import static java.time.temporal.ChronoField.ALIGNED_WEEK_OF_YEAR;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.opensearch.sql.data.type.ExprCoreType.LONG;
 
+import java.time.Instant;
 import java.time.LocalDate;
+import java.time.LocalDateTime;
+import java.time.LocalTime;
+import java.time.ZoneId;
+import java.time.temporal.IsoFields;
 import java.util.stream.Stream;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.MethodSource;
+import org.junit.jupiter.params.provider.ValueSource;
 import org.opensearch.sql.data.model.ExprDateValue;
 import org.opensearch.sql.data.model.ExprTimeValue;
 import org.opensearch.sql.data.model.ExprTimestampValue;
@@ -23,6 +28,7 @@
 import org.opensearch.sql.expression.Expression;
 import org.opensearch.sql.expression.ExpressionTestBase;
 import org.opensearch.sql.expression.FunctionExpression;
+import org.opensearch.sql.expression.function.FunctionProperties;
 
 class ExtractTest extends ExpressionTestBase {
 
@@ -82,9 +88,14 @@ public void testExtractWithDatetime(String part, long expected) {
   }
 
   private void datePartWithTimeArgQuery(String part, String time, long expected) {
+    datePartWithTimeArgQuery(functionProperties, part, time, expected);
+  }
+
+  private void datePartWithTimeArgQuery(
+      FunctionProperties properties, String part, String time, long expected) {
     ExprTimeValue timeValue = new ExprTimeValue(time);
     FunctionExpression datetimeExpression =
-        DSL.extract(functionProperties, DSL.literal(part), DSL.literal(timeValue));
+        DSL.extract(properties, DSL.literal(part), DSL.literal(timeValue));
 
     assertEquals(LONG, datetimeExpression.type());
     assertEquals(expected, eval(datetimeExpression).longValue());
@@ -96,21 +107,49 @@ public void testExtractDatePartWithTimeType() {
 
     datePartWithTimeArgQuery("DAY", timeInput, now.getDayOfMonth());
 
-    // To avoid flaky test, skip the testing in December and January because the WEEK is ISO 8601
-    // week-of-week-based-year which is considered to start on a Monday and week 1 is the first week
-    // with >3 days. it is possible for early-January dates to be part of the 52nd or 53rd week of
-    // the previous year, and for late-December dates to be part of the first week of the next year.
-    // For example, 2005-01-02 is part of the 53rd week of year 2004, while 2012-12-31 is part of
-    // the first week of 2013
-    if (now.getMonthValue() != 1 && now.getMonthValue() != 12) {
-      datePartWithTimeArgQuery("WEEK", datetimeInput, now.get(ALIGNED_WEEK_OF_YEAR));
-    }
-
     datePartWithTimeArgQuery("MONTH", timeInput, now.getMonthValue());
 
     datePartWithTimeArgQuery("YEAR", timeInput, now.getYear());
   }
 
+  @ParameterizedTest(name = "{0}")
+  @ValueSource(
+      strings = {
+        "2009-12-26",
+        "2009-12-27",
+        "2008-12-28", // Week 52 of week-based-year 2008
+        "2009-12-29",
+        "2008-12-29", // Week 1 of week-based-year 2009
+        "2008-12-31", // Week 1 of week-based-year 2009
+        "2009-01-01", // Week 1 of week-based-year 2009
+        "2009-01-04", // Week 1 of week-based-year 2009
+        "2009-01-05", // Week 2 of week-based-year 2009
+        "2025-12-27", //  year with 52 weeks
+        "2026-01-01", // year starts on a THURSDAY
+        "2028-12-30", //  year with 53 weeks
+        "2028-12-31", // year starts in December
+        "2029-01-01",
+        "2033-12-31", // year with 53 weeks
+        "2034-01-01", // January 1st on a SUNDAY
+        "2034-12-30", // year with 52 weeks
+        "2034-12-31"
+      })
+  public void testExtractWeekPartWithTimeType(String arg) {
+
+    // setup default date/time properties for the extract function
+    ZoneId currentZoneId = ZoneId.systemDefault();
+    Instant nowInstant =
+        LocalDate.parse(arg).atTime(LocalTime.parse(timeInput)).atZone(currentZoneId).toInstant();
+    FunctionProperties properties = new FunctionProperties(nowInstant, currentZoneId);
+
+    // Expected WEEK value should be formated from week-of-week-based-year
+    LocalDateTime localDateTime = LocalDateTime.ofInstant(nowInstant, currentZoneId);
+    int expected = localDateTime.get(IsoFields.WEEK_OF_WEEK_BASED_YEAR);
+
+    // verify
+    datePartWithTimeArgQuery(properties, "WEEK", timeInput, expected);
+  }
+
   @ParameterizedTest(name = "{0}")
   @MethodSource("getDateResultsForExtractFunction")
   public void testExtractWithDate(String part, long expected) {

core/src/test/java/org/opensearch/sql/expression/function/OpenSearchFunctionsTest.java

@@ -22,6 +22,7 @@
 import org.opensearch.sql.expression.ExpressionTestBase;
 import org.opensearch.sql.expression.FunctionExpression;
 import org.opensearch.sql.expression.NamedArgumentExpression;
+import org.opensearch.sql.expression.ReferenceExpression;
 import org.opensearch.sql.expression.env.Environment;
 
 public class OpenSearchFunctionsTest extends ExpressionTestBase {
@@ -309,4 +310,17 @@ void nested_query() {
     assertEquals(expr.valueOf(nestedTuple), ExprValueUtils.stringValue("result"));
     assertEquals(expr.type(), STRING);
   }
+
+  @Test
+  void opensearchExecutableFunction_valueOf() {
+    var ipInStr =
+        new OpenSearchFunctions.OpenSearchExecutableFunction(
+            BuiltinFunctionName.GEOIP.getName(),
+            List.of(DSL.literal("my-datasource"), new ReferenceExpression("ipInStr", STRING)),
+            BOOLEAN);
+    assertThrows(
+        UnsupportedOperationException.class,
+        () -> ipInStr.valueOf(valueEnv()),
+        "OpenSearch defined function [geoip] is only supported in Eval operation.");
+  }
 }