Support decimal literal in span function

Signed-off-by: Yuanchun Shen <yuanchu@amazon.com>

Author: yuancu

core/src/main/java/org/opensearch/sql/expression/function/udf/SpanFunction.java

@@ -20,7 +20,6 @@
 import org.apache.calcite.schema.impl.ScalarFunctionImpl;
 import org.apache.calcite.sql.type.CompositeOperandTypeChecker;
 import org.apache.calcite.sql.type.OperandTypes;
-import org.apache.calcite.sql.type.ReturnTypes;
 import org.apache.calcite.sql.type.SqlReturnTypeInference;
 import org.apache.calcite.sql.type.SqlTypeFamily;
 import org.apache.calcite.sql.type.SqlTypeUtil;
@@ -44,7 +43,17 @@ public SpanFunction() {
 
   @Override
   public SqlReturnTypeInference getReturnTypeInference() {
-    return ReturnTypes.ARG0;
+    // Return arg0 type if it has a unit (i.e. time related span)
+    return callBinding -> {
+      if (SqlTypeUtil.isString(callBinding.getOperandType(2))) {
+        return callBinding.getOperandType(0);
+      }
+      // Use the least restrictive type between the field type and the interval type if it's a
+      // numeric span. E.g. span(int_field, double_literal) -> double
+      return callBinding
+          .getTypeFactory()
+          .leastRestrictive(List.of(callBinding.getOperandType(0), callBinding.getOperandType(1)));
+    };
   }
 
   @Override
@@ -56,10 +65,9 @@ public UDFOperandMetadata getOperandMetadata() {
                 .or(
                     OperandTypes.family(
                         SqlTypeFamily.DATETIME, SqlTypeFamily.NUMERIC, SqlTypeFamily.CHARACTER))
-                // TODO: numeric span should support decimal as its interval
                 .or(
                     OperandTypes.family(
-                        SqlTypeFamily.NUMERIC, SqlTypeFamily.INTEGER, SqlTypeFamily.ANY)));
+                        SqlTypeFamily.NUMERIC, SqlTypeFamily.NUMERIC, SqlTypeFamily.ANY)));
   }
 
   public static class SpanImplementor implements NotNullImplementor {

integ-test/src/test/java/org/opensearch/sql/ppl/StatsCommandIT.java

@@ -806,4 +806,21 @@ public void testStatsByDependentGroupFields() throws IOException {
     verifyDataRows(
         response, rows(61, 310, 41, 10, 31), rows(60, 390, 49, 10, 39), rows(59, 260, 36, 10, 26));
   }
+
+  @Test
+  public void testStatsByFractionalSpan() throws IOException {
+    JSONObject response1 =
+        executeQuery(
+            String.format(
+                "source=%s | stats count by span(balance, 4170.5)",
+                TEST_INDEX_BANK_WITH_NULL_VALUES));
+    verifySchema(response1, schema("count", "bigint"), schema("span(balance,4170.5)", "double"));
+    verifyDataRows(
+        response1,
+        rows(3, null),
+        rows(1, 4170.5),
+        rows(1, 29193.5),
+        rows(1, 37534.5),
+        rows(1, 45875.5));
+  }
 }

ppl/src/test/java/org/opensearch/sql/ppl/parser/AstBuilderTest.java

@@ -1488,4 +1488,18 @@ public void testReplaceCommandWithMultiplePairs() {
     // Test multiple pattern/replacement pairs
     plan("source=t | replace 'a' WITH 'A', 'b' WITH 'B' IN field");
   }
+
+  @Test
+  public void testTimeSpanWithDecimalShouldThrow() {
+    Throwable t1 =
+        assertThrows(
+            SyntaxCheckException.class, () -> plan("source=t | timechart  span=1.5d count"));
+    assertTrue(t1.getMessage().contains("[1.5d] is not a valid term at this part of the query"));
+
+    Throwable t2 =
+        assertThrows(
+            SyntaxCheckException.class,
+            () -> plan("source=t | stats count by span(@timestamp, 2.5y)"));
+    assertTrue(t2.getMessage().contains("[y] is not a valid term at this part of the query"));
+  }
 }