Jackie-Jiang commented on code in PR #9184: URL: https://github.com/apache/pinot/pull/9184#discussion_r953182276
########## pinot-common/src/test/java/org/apache/pinot/sql/parsers/CalciteSqlCompilerTest.java: ########## @@ -184,6 +184,44 @@ public void testQuotedStrings() { .getIdentifier().getName(), "Martha''s Vineyard"); } + @Test + public void testExtract() { + try { Review Comment: We can remove the try-catch and add `throw Exception` to the test method ########## pinot-common/src/test/java/org/apache/pinot/sql/parsers/CalciteSqlCompilerTest.java: ########## @@ -184,6 +184,44 @@ public void testQuotedStrings() { .getIdentifier().getName(), "Martha''s Vineyard"); } + @Test + public void testExtract() { + try { + //@formatter:off Review Comment: I don't think we need to turn off/on the formatter since there is no special formatting needed ########## pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/ExtractTransformFunction.java: ########## @@ -0,0 +1,97 @@ +package org.apache.pinot.core.operator.transform.function; + +import java.util.List; +import java.util.Map; +import org.apache.pinot.core.operator.blocks.ProjectionBlock; +import org.apache.pinot.core.operator.transform.TransformResultMetadata; +import org.apache.pinot.segment.spi.datasource.DataSource; +import org.joda.time.Chronology; +import org.joda.time.DateTimeField; +import org.joda.time.chrono.ISOChronology; + + +public class ExtractTransformFunction extends BaseTransformFunction { + public static final String FUNCTION_NAME = "extract"; + private TransformFunction _mainTransformFunction; + protected String _field; + protected Chronology _chronology = ISOChronology.getInstanceUTC(); + + @Override + public String getName() { return FUNCTION_NAME; } + + @Override + public void init(List<TransformFunction> arguments, Map<String, DataSource> dataSourceMap) { + if (arguments.size() != 2) { + throw new IllegalArgumentException("Exactly 2 arguments are required for EXTRACT transform function"); + } + + _field = ((LiteralTransformFunction) arguments.get(0)).getLiteral(); + if(!validField(_field)) { + throw new IllegalArgumentException("FIELD in EXTRACT expression is not valid"); + } + + _mainTransformFunction = arguments.get(1); + } + + private boolean validField(String field) { Review Comment: Suggest making a `enum` for the `Field`, then use `Field.valueOf()` and directly store the `Field` in the member variable ########## pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/ExtractTransformFunction.java: ########## @@ -0,0 +1,97 @@ +package org.apache.pinot.core.operator.transform.function; + +import java.util.List; +import java.util.Map; +import org.apache.pinot.core.operator.blocks.ProjectionBlock; +import org.apache.pinot.core.operator.transform.TransformResultMetadata; +import org.apache.pinot.segment.spi.datasource.DataSource; +import org.joda.time.Chronology; +import org.joda.time.DateTimeField; +import org.joda.time.chrono.ISOChronology; + + +public class ExtractTransformFunction extends BaseTransformFunction { + public static final String FUNCTION_NAME = "extract"; + private TransformFunction _mainTransformFunction; + protected String _field; + protected Chronology _chronology = ISOChronology.getInstanceUTC(); + + @Override + public String getName() { return FUNCTION_NAME; } + + @Override + public void init(List<TransformFunction> arguments, Map<String, DataSource> dataSourceMap) { + if (arguments.size() != 2) { + throw new IllegalArgumentException("Exactly 2 arguments are required for EXTRACT transform function"); + } + + _field = ((LiteralTransformFunction) arguments.get(0)).getLiteral(); + if(!validField(_field)) { Review Comment: (code format) Please reformat using the [Pinot Style](https://docs.pinot.apache.org/developers/developers-and-contributors/code-setup#intellij) ########## pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/ExtractTransformFunction.java: ########## @@ -0,0 +1,97 @@ +package org.apache.pinot.core.operator.transform.function; + +import java.util.List; +import java.util.Map; +import org.apache.pinot.core.operator.blocks.ProjectionBlock; +import org.apache.pinot.core.operator.transform.TransformResultMetadata; +import org.apache.pinot.segment.spi.datasource.DataSource; +import org.joda.time.Chronology; +import org.joda.time.DateTimeField; +import org.joda.time.chrono.ISOChronology; + + +public class ExtractTransformFunction extends BaseTransformFunction { + public static final String FUNCTION_NAME = "extract"; + private TransformFunction _mainTransformFunction; + protected String _field; + protected Chronology _chronology = ISOChronology.getInstanceUTC(); + + @Override + public String getName() { return FUNCTION_NAME; } + + @Override + public void init(List<TransformFunction> arguments, Map<String, DataSource> dataSourceMap) { + if (arguments.size() != 2) { + throw new IllegalArgumentException("Exactly 2 arguments are required for EXTRACT transform function"); + } + + _field = ((LiteralTransformFunction) arguments.get(0)).getLiteral(); + if(!validField(_field)) { + throw new IllegalArgumentException("FIELD in EXTRACT expression is not valid"); + } + + _mainTransformFunction = arguments.get(1); + } + + private boolean validField(String field) { + String[] fieldsList = {"YEAR", "MONTH", "DAY", "HOUR", "MINUTE", "SECOND"}; + for(String it : fieldsList) { + if(it.equals(field)) { return true; } + } + return false; + } + + @Override + public TransformResultMetadata getResultMetadata() { + return INT_SV_NO_DICTIONARY_METADATA; + } + + @Override + public int[] transformToIntValuesSV(ProjectionBlock projectionBlock) { + int numDocs = projectionBlock.getNumDocs(); + + if(_intValuesSV == null || _intValuesSV.length < numDocs) { + _intValuesSV = new int[numDocs]; + } + + long[] timestamps = _mainTransformFunction.transformToLongValuesSV(projectionBlock); + + convert(timestamps, numDocs, _intValuesSV); + + 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 "MONTH": + accessor = _chronology.monthOfYear(); + output[i] = accessor.get(timestamps[i]); + break; + case "DAY": + accessor = _chronology.dayOfMonth(); + 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: Review Comment: throw exception here in case we missed some unhandled case -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: commits-unsubscr...@pinot.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: commits-unsubscr...@pinot.apache.org For additional commands, e-mail: commits-h...@pinot.apache.org