This is an automated email from the ASF dual-hosted git repository.
jackie pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new a0e861e36f Support 'EXTRACT' as a scalar function (#13463)
a0e861e36f is described below
commit a0e861e36f198f64b93cfcae81b95a2f7e8a9703
Author: Yash Mayya <[email protected]>
AuthorDate: Sat Jun 29 05:14:16 2024 +0530
Support 'EXTRACT' as a scalar function (#13463)
---
.../pinot/common/function/DateTimeUtils.java | 40 ++++++++++++-
.../common/function/scalar/DateTimeFunctions.java | 5 ++
.../pinot/sql/parsers/CalciteSqlCompilerTest.java | 37 +++++-------
.../function/ExtractTransformFunction.java | 66 ++--------------------
.../function/ExtractTransformFunctionTest.java | 3 +
.../rel/rules/PinotEvaluateLiteralRule.java | 3 +
6 files changed, 70 insertions(+), 84 deletions(-)
diff --git
a/pinot-common/src/main/java/org/apache/pinot/common/function/DateTimeUtils.java
b/pinot-common/src/main/java/org/apache/pinot/common/function/DateTimeUtils.java
index b2a8ce933c..f211bc71f3 100644
---
a/pinot-common/src/main/java/org/apache/pinot/common/function/DateTimeUtils.java
+++
b/pinot-common/src/main/java/org/apache/pinot/common/function/DateTimeUtils.java
@@ -33,7 +33,7 @@ import org.slf4j.LoggerFactory;
/**
- * Helper methods and constructs for datetrunc function
+ * Helper methods and constructs for date/time functions
*/
public class DateTimeUtils {
private DateTimeUtils() {
@@ -41,6 +41,7 @@ public class DateTimeUtils {
private static final Logger LOGGER =
LoggerFactory.getLogger(DateTimeUtils.class);
private static final DateTimeFieldType QUARTER_OF_YEAR = new
QuarterOfYearDateTimeField();
+ private static final Chronology CHRONOLOGY_UTC =
ISOChronology.getInstanceUTC();
public static DateTimeField getTimestampField(ISOChronology chronology,
String unitString) {
switch (unitString.toLowerCase()) {
@@ -152,4 +153,41 @@ public class DateTimeUtils {
}
}
}
+
+ /**
+ * The supported field types for the EXTRACT operator
+ */
+ public enum ExtractFieldType {
+ YEAR, QUARTER, MONTH, WEEK, DAY, DOY, DOW, HOUR, MINUTE, SECOND
+ }
+
+ /**
+ * Helper method to implement the SQL <code>EXTRACT</code> operator.
+ */
+ public static int extract(ExtractFieldType extractFieldType, long timestamp)
{
+ switch (extractFieldType) {
+ case YEAR:
+ return CHRONOLOGY_UTC.year().get(timestamp);
+ case QUARTER:
+ return (CHRONOLOGY_UTC.monthOfYear().get(timestamp) - 1) / 3 + 1;
+ case MONTH:
+ return CHRONOLOGY_UTC.monthOfYear().get(timestamp);
+ case WEEK:
+ return CHRONOLOGY_UTC.weekOfWeekyear().get(timestamp);
+ case DAY:
+ return CHRONOLOGY_UTC.dayOfMonth().get(timestamp);
+ case DOY:
+ return CHRONOLOGY_UTC.dayOfYear().get(timestamp);
+ case DOW:
+ return CHRONOLOGY_UTC.dayOfWeek().get(timestamp);
+ case HOUR:
+ return CHRONOLOGY_UTC.hourOfDay().get(timestamp);
+ case MINUTE:
+ return CHRONOLOGY_UTC.minuteOfHour().get(timestamp);
+ case SECOND:
+ return CHRONOLOGY_UTC.secondOfMinute().get(timestamp);
+ default:
+ throw new IllegalArgumentException("Unsupported FIELD type");
+ }
+ }
}
diff --git
a/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/DateTimeFunctions.java
b/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/DateTimeFunctions.java
index adb6c9bda6..ac734dd4da 100644
---
a/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/DateTimeFunctions.java
+++
b/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/DateTimeFunctions.java
@@ -1259,4 +1259,9 @@ public class DateTimeFunctions {
}
return results;
}
+
+ @ScalarFunction
+ public static int extract(String interval, long timestamp) {
+ return
DateTimeUtils.extract(DateTimeUtils.ExtractFieldType.valueOf(interval),
timestamp);
+ }
}
diff --git
a/pinot-common/src/test/java/org/apache/pinot/sql/parsers/CalciteSqlCompilerTest.java
b/pinot-common/src/test/java/org/apache/pinot/sql/parsers/CalciteSqlCompilerTest.java
index 3576c19660..7fb1f7ae10 100644
---
a/pinot-common/src/test/java/org/apache/pinot/sql/parsers/CalciteSqlCompilerTest.java
+++
b/pinot-common/src/test/java/org/apache/pinot/sql/parsers/CalciteSqlCompilerTest.java
@@ -263,32 +263,25 @@ public class CalciteSqlCompilerTest {
@Test
public void testExtract() {
{
- // Case 1 -- Year and date format ('2017-06-15')
- PinotQuery pinotQuery = compileToPinotQuery("SELECT EXTRACT(YEAR FROM
'2017-06-15')");
- Function function = pinotQuery.getSelectList().get(0).getFunctionCall();
-
Assert.assertEquals(function.getOperands().get(0).getLiteral().getStringValue(),
"YEAR");
-
Assert.assertEquals(function.getOperands().get(1).getLiteral().getStringValue(),
"2017-06-15");
+ // Case 1 -- Year
+ PinotQuery pinotQuery = compileToPinotQuery("SELECT EXTRACT(YEAR FROM
1719573611000)");
+ // The CompileTimeFunctionsInvoker will rewrite the query to replace the
function call with the resultant literal
+ // value
+
Assert.assertEquals(pinotQuery.getSelectList().get(0).getLiteral().getIntValue(),
2024);
}
{
- // Case 2 -- date format ('2017-06-15 09:34:21')
- PinotQuery pinotQuery = compileToPinotQuery("SELECT EXTRACT(YEAR FROM
'2017-06-15 09:34:21')");
- Function function = pinotQuery.getSelectList().get(0).getFunctionCall();
-
Assert.assertEquals(function.getOperands().get(0).getLiteral().getStringValue(),
"YEAR");
-
Assert.assertEquals(function.getOperands().get(1).getLiteral().getStringValue(),
"2017-06-15 09:34:21");
+ // Case 2 -- Month
+ PinotQuery pinotQuery = compileToPinotQuery("SELECT EXTRACT(MONTH FROM
'1719573611000')");
+ // The CompileTimeFunctionsInvoker will rewrite the query to replace the
function call with the resultant literal
+ // value
+
Assert.assertEquals(pinotQuery.getSelectList().get(0).getLiteral().getIntValue(),
6);
}
{
- // Case 3 -- Month
- PinotQuery pinotQuery = compileToPinotQuery("SELECT EXTRACT(MONTH FROM
'2017-06-15')");
- Function function = pinotQuery.getSelectList().get(0).getFunctionCall();
-
Assert.assertEquals(function.getOperands().get(0).getLiteral().getStringValue(),
"MONTH");
-
Assert.assertEquals(function.getOperands().get(1).getLiteral().getStringValue(),
"2017-06-15");
- }
- {
- // Case 4 -- Day
- PinotQuery pinotQuery = compileToPinotQuery("SELECT EXTRACT(DAY FROM
'2017-06-15')");
- Function function = pinotQuery.getSelectList().get(0).getFunctionCall();
-
Assert.assertEquals(function.getOperands().get(0).getLiteral().getStringValue(),
"DAY");
-
Assert.assertEquals(function.getOperands().get(1).getLiteral().getStringValue(),
"2017-06-15");
+ // Case 3 -- Day
+ PinotQuery pinotQuery = compileToPinotQuery("SELECT EXTRACT(DAY FROM
1719573611000)");
+ // The CompileTimeFunctionsInvoker will rewrite the query to replace the
function call with the resultant literal
+ // value
+
Assert.assertEquals(pinotQuery.getSelectList().get(0).getLiteral().getIntValue(),
28);
}
}
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/ExtractTransformFunction.java
b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/ExtractTransformFunction.java
index 7bbc19acde..906ce38781 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/ExtractTransformFunction.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/ExtractTransformFunction.java
@@ -20,24 +20,18 @@ package org.apache.pinot.core.operator.transform.function;
import java.util.List;
import java.util.Map;
+import java.util.stream.IntStream;
+import org.apache.pinot.common.function.DateTimeUtils;
import org.apache.pinot.core.operator.ColumnContext;
import org.apache.pinot.core.operator.blocks.ValueBlock;
import org.apache.pinot.core.operator.transform.TransformResultMetadata;
-import org.joda.time.Chronology;
-import org.joda.time.DateTimeField;
-import org.joda.time.chrono.ISOChronology;
import org.roaringbitmap.RoaringBitmap;
public class ExtractTransformFunction extends BaseTransformFunction {
public static final String FUNCTION_NAME = "extract";
private TransformFunction _mainTransformFunction;
- protected Field _field;
- protected Chronology _chronology = ISOChronology.getInstanceUTC();
-
- private enum Field {
- YEAR, QUARTER, MONTH, WEEK, DAY, DOY, DOW, HOUR, MINUTE, SECOND
- }
+ private DateTimeUtils.ExtractFieldType _field;
@Override
public String getName() {
@@ -51,7 +45,7 @@ public class ExtractTransformFunction extends
BaseTransformFunction {
throw new IllegalArgumentException("Exactly 2 arguments are required for
EXTRACT transform function");
}
- _field = Field.valueOf(((LiteralTransformFunction)
arguments.get(0)).getStringLiteral());
+ _field =
DateTimeUtils.ExtractFieldType.valueOf(((LiteralTransformFunction)
arguments.get(0)).getStringLiteral());
_mainTransformFunction = arguments.get(1);
}
@@ -65,60 +59,10 @@ public class ExtractTransformFunction extends
BaseTransformFunction {
int numDocs = valueBlock.getNumDocs();
initIntValuesSV(numDocs);
long[] timestamps =
_mainTransformFunction.transformToLongValuesSV(valueBlock);
- convert(timestamps, numDocs, _intValuesSV);
+ IntStream.range(0, numDocs).forEach(i -> _intValuesSV[i] =
DateTimeUtils.extract(_field, timestamps[i]));
return _intValuesSV;
}
- private void convert(long[] timestamps, int numDocs, int[] output) {
- for (int i = 0; i < numDocs; i++) {
- DateTimeField accessor;
- switch (_field) {
- case YEAR:
- accessor = _chronology.year();
- output[i] = accessor.get(timestamps[i]);
- break;
- case QUARTER:
- accessor = _chronology.monthOfYear();
- output[i] = (accessor.get(timestamps[i]) - 1) / 3 + 1;
- break;
- case MONTH:
- accessor = _chronology.monthOfYear();
- output[i] = accessor.get(timestamps[i]);
- break;
- case WEEK:
- accessor = _chronology.weekOfWeekyear();
- output[i] = accessor.get(timestamps[i]);
- break;
- case DAY:
- accessor = _chronology.dayOfMonth();
- output[i] = accessor.get(timestamps[i]);
- break;
- case DOY:
- accessor = _chronology.dayOfYear();
- output[i] = accessor.get(timestamps[i]);
- break;
- case DOW:
- accessor = _chronology.dayOfWeek();
- output[i] = accessor.get(timestamps[i]);
- break;
- case HOUR:
- accessor = _chronology.hourOfDay();
- output[i] = accessor.get(timestamps[i]);
- break;
- case MINUTE:
- accessor = _chronology.minuteOfHour();
- output[i] = accessor.get(timestamps[i]);
- break;
- case SECOND:
- accessor = _chronology.secondOfMinute();
- output[i] = accessor.get(timestamps[i]);
- break;
- default:
- throw new IllegalArgumentException("Unsupported FIELD type");
- }
- }
- }
-
@Override
public RoaringBitmap getNullBitmap(ValueBlock valueBlock) {
return _mainTransformFunction.getNullBitmap(valueBlock);
diff --git
a/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/ExtractTransformFunctionTest.java
b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/ExtractTransformFunctionTest.java
index ea74f266c5..96eda1c96c 100644
---
a/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/ExtractTransformFunctionTest.java
+++
b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/ExtractTransformFunctionTest.java
@@ -38,9 +38,12 @@ public class ExtractTransformFunctionTest extends
BaseTransformFunctionTest {
return new Object[][]{
//@formatter:off
{"year", (LongToIntFunction) DateTimeFunctions::year},
+ {"quarter", (LongToIntFunction) timestamp ->
(DateTimeFunctions.monthOfYear(timestamp) - 1) / 3 + 1},
{"month", (LongToIntFunction) DateTimeFunctions::monthOfYear},
{"week", (LongToIntFunction) DateTimeFunctions::weekOfYear},
{"day", (LongToIntFunction) DateTimeFunctions::dayOfMonth},
+ {"doy", (LongToIntFunction) DateTimeFunctions::dayOfYear},
+ {"dow", (LongToIntFunction) DateTimeFunctions::dayOfWeek},
{"hour", (LongToIntFunction) DateTimeFunctions::hour},
{"minute", (LongToIntFunction) DateTimeFunctions::minute},
{"second", (LongToIntFunction) DateTimeFunctions::second},
diff --git
a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotEvaluateLiteralRule.java
b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotEvaluateLiteralRule.java
index 5ea790b408..c0b7b4b21c 100644
---
a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotEvaluateLiteralRule.java
+++
b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotEvaluateLiteralRule.java
@@ -27,6 +27,7 @@ import java.util.List;
import java.util.Objects;
import javax.annotation.Nullable;
import org.apache.calcite.avatica.util.ByteString;
+import org.apache.calcite.avatica.util.TimeUnitRange;
import org.apache.calcite.plan.RelOptRule;
import org.apache.calcite.plan.RelOptRuleCall;
import org.apache.calcite.rel.logical.LogicalFilter;
@@ -233,6 +234,8 @@ public class PinotEvaluateLiteralRule {
} else if (value instanceof ByteString) {
// BYTES
return ((ByteString) value).getBytes();
+ } else if (value instanceof TimeUnitRange) {
+ return ((TimeUnitRange) value).name();
} else {
return value;
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
