Skip to content

Add support for arrayConcatLong, arrayConcatFloat, arrayConcatDouble #9131

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
siddharthteotia merged 1 commit into from
Aug 9, 2022
Merged

Add support for arrayConcatLong, arrayConcatFloat, arrayConcatDouble #9131

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 15 additions & 0 deletions pinot-common/src/main/java/org/apache/pinot/common/function/scalar/ArrayFunctions.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -132,6 +132,21 @@ public static int[] arrayConcatInt(int[] values1, int[] values2) {
return ArrayUtils.addAll(values1, values2);
}

@ScalarFunction
public static long[] arrayConcatLong(long[] values1, long[] values2) {
return ArrayUtils.addAll(values1, values2);
}

@ScalarFunction
public static float[] arrayConcatFloat(float[] values1, float[] values2) {
return ArrayUtils.addAll(values1, values2);
}

@ScalarFunction
public static double[] arrayConcatDouble(double[] values1, double[] values2) {
return ArrayUtils.addAll(values1, values2);
}

@ScalarFunction
public static String[] arrayConcatString(String[] values1, String[] values2) {
return ArrayUtils.addAll(values1, values2);
Expand Down
51 changes: 51 additions & 0 deletions ...org/apache/pinot/core/operator/transform/function/ScalarTransformFunctionWrapperTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -803,6 +803,57 @@ public void testArrayConcatIntTransformFunction() {
testTransformFunctionMV(transformFunction, expectedValues);
}

@Test
Copy link
Contributor

@siddharthteotia siddharthteotia Jul 29, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please also consider adding some other query execution tests. Take a look at BaseMultiValueQueriesTest or OfflineClusterIntegrationTest to add such tests.

Example cases ...

Nest the function

SELECT arraylength(arrayConcatInt(mvIntCol1, mvIntCol2))
FROM FOO

SELECT arraySortInt(arrayConcatInt(mvIntCol1, mvIntCol2))
FROM FOO

Use the output of the expression in WHERE clause

SELECT arrayConcatInt(mvIntCol1, mvIntCol2) AS concatenated
FROM FOO
WHERE arrayLength(concatenated) >= 2

Use in GROUP BY

SELECT arrayConcatInt(mvIntCol1, mvIntCol2) AS concatenated, sum(intcol)
FROM FOO
WHERE arrayLength(concatenated) >= 2
GROUP BY concatenated

Copy link
Contributor Author

@jugomezv jugomezv Jul 29, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sidd:

Sure this is a good idea, just to confirm: such tests do not exists for the currently supported functions like arrayConcatInt and arrayConcatString, right?

If that is the case I can look into adding these tests for the new functions being added in this PR, I can file a new issue to support those that were added earlier so we don't hold this on earlier tech debt.

Copy link
Contributor

@siddharthteotia siddharthteotia Jul 29, 2022 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

such tests do not exists for the currently supported functions like arrayConcatInt and arrayConcatString, right?

I think that is correct (most likely). But we will have to dig through since tests are spread across so many files unfortunately.

If that is the case I can look into adding these tests for the new functions being added in this PR

Yes, in the PR I suggest adding tests only for new functions. Separate PR can be filed to add more tests for existing functions.

public void testArrayConcatLongTransformFunction() {
ExpressionContext expression = RequestContextUtils.getExpression(
String.format("array_concat_long(%s, %s)", LONG_MV_COLUMN, LONG_MV_COLUMN));
TransformFunction transformFunction = TransformFunctionFactory.get(expression, _dataSourceMap);
assertTrue(transformFunction instanceof ScalarTransformFunctionWrapper);
assertEquals(transformFunction.getName(), "arrayConcatLong");
assertEquals(transformFunction.getResultMetadata().getDataType(), DataType.LONG);
assertFalse(transformFunction.getResultMetadata().isSingleValue());
long[][] expectedValues = new long[NUM_ROWS][];
for (int i = 0; i < NUM_ROWS; i++) {
expectedValues[i] = _longMVValues[i].clone();
expectedValues[i] = ArrayUtils.addAll(expectedValues[i], expectedValues[i]);
}
testTransformFunctionMV(transformFunction, expectedValues);
}

@Test
public void testArrayConcatFloatTransformFunction() {
ExpressionContext expression = RequestContextUtils.getExpression(
String.format("array_concat_float(%s, %s)", FLOAT_MV_COLUMN, FLOAT_MV_COLUMN));
TransformFunction transformFunction = TransformFunctionFactory.get(expression, _dataSourceMap);
assertTrue(transformFunction instanceof ScalarTransformFunctionWrapper);
assertEquals(transformFunction.getName(), "arrayConcatFloat");
assertEquals(transformFunction.getResultMetadata().getDataType(), DataType.FLOAT);
assertFalse(transformFunction.getResultMetadata().isSingleValue());
float[][] expectedValues = new float[NUM_ROWS][];
for (int i = 0; i < NUM_ROWS; i++) {
expectedValues[i] = _floatMVValues[i].clone();
expectedValues[i] = ArrayUtils.addAll(expectedValues[i], expectedValues[i]);
}
testTransformFunctionMV(transformFunction, expectedValues);
}

@Test
public void testArrayConcatDoubleTransformFunction() {
ExpressionContext expression = RequestContextUtils.getExpression(
String.format("array_concat_double(%s, %s)", DOUBLE_MV_COLUMN, DOUBLE_MV_COLUMN));
TransformFunction transformFunction = TransformFunctionFactory.get(expression, _dataSourceMap);
assertTrue(transformFunction instanceof ScalarTransformFunctionWrapper);
assertEquals(transformFunction.getName(), "arrayConcatDouble");
assertEquals(transformFunction.getResultMetadata().getDataType(), DataType.DOUBLE);
assertFalse(transformFunction.getResultMetadata().isSingleValue());
double[][] expectedValues = new double[NUM_ROWS][];
for (int i = 0; i < NUM_ROWS; i++) {
expectedValues[i] = _doubleMVValues[i].clone();
expectedValues[i] = ArrayUtils.addAll(expectedValues[i], expectedValues[i]);
}
testTransformFunctionMV(transformFunction, expectedValues);
}

@Test
public void testConcatStringTransformFunction() {
ExpressionContext expression = RequestContextUtils.getExpression(
Expand Down
299 changes: 299 additions & 0 deletions pinot-core/src/test/java/org/apache/pinot/queries/MultiValueRawQueriesTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,7 @@ public class MultiValueRawQueriesTest extends BaseQueriesTest {
private static final int NUM_UNIQUE_RECORDS_PER_SEGMENT = 10;
private static final int NUM_DUPLICATES_PER_RECORDS = 2;
private static final int MV_OFFSET = 100;
private static final int MV_LENGTH = 2; // Entries in each multivalue array
private static final int BASE_VALUE_1 = 0;
private static final int BASE_VALUE_2 = 1000;

Expand Down Expand Up @@ -2045,4 +2046,302 @@ public void testTransformInsideAggregateQueries() {
assertEquals(value[0], 8L);
}
}

private boolean validateMvArrayLength(String query, int expectedLength, int expectedRecords) {
ResultTable resultTable = getBrokerResponse(query).getResultTable();
assertNotNull(resultTable);
List<Object[]> recordRows = resultTable.getRows();
assertNotNull(recordRows);
assertEquals(recordRows.size(), expectedRecords);
for (int i = 0; i < recordRows.size(); i++) {
Object[] values = recordRows.get(i);
assertNotNull(values);
int len = (int) values[0];
if (len != expectedLength) {
return false;
}
}
return true;
}

private boolean validateQueryOutputSize(String query, int expectedRecords) {
ResultTable resultTable = getBrokerResponse(query).getResultTable();
assertNotNull(resultTable);
List<Object[]> recordRows = resultTable.getRows();
assertNotNull(recordRows);
assertEquals(recordRows.size(), expectedRecords);
return true;
}

private String getConcatQueryString(String concatFunction, String col1, String col2, String concatCol, String table,
int limit) {
return String.format("SELECT %s(%s, %s) AS %s FROM %s LIMIT %d", concatFunction, col1, col2, concatCol, table,
limit);
}

private String getConcatNestingQueryString(String concatFunction, String col1, String col2, String concatCol,
String table, int limit) {
return String.format("SELECT arraylength(%s(%s, %s)) AS %s FROM %s LIMIT %d", concatFunction, col1, col2, concatCol,
table, limit);
}

private String getConcatUseInWhereQueryString(String concatFunction, String col1, String col2, String concatCol,
String table, String compareOperator, int mvPerArray, int limit) {
return String.format("SELECT %s(%s, %s) AS %s FROM %s WHERE arraylength(%s) %s %d LIMIT %d", concatFunction, col1,
col2, concatCol, table, concatCol, compareOperator, mvPerArray, limit);
}

private String getConcatGroupByQueryString(String concatFunction, String col1, String col2, String concatCol,
String table, String compareOperator, int mvPerArray, int limit) {
return String.format(
"SELECT %s(%s, %s) AS %s, sum(svIntCol) FROM %s WHERE arraylength(%s) %s %d GROUP BY %s LIMIT %d",
concatFunction, col1, col2, concatCol, table, concatCol, compareOperator, mvPerArray, concatCol, limit);
}

@Test
public void testArrayConcatInt() {
{
// Test basic arrayConcatInt() functionality
String query = getConcatQueryString("arrayConcatInt", "mvIntCol", "mvRawIntCol", "concatInts", "testTable",
NUM_UNIQUE_RECORDS_PER_SEGMENT);
ResultTable resultTable = getBrokerResponse(query).getResultTable();
assertNotNull(resultTable);
List<Object[]> recordRows = resultTable.getRows();
assertNotNull(recordRows);
assertEquals(recordRows.size(), NUM_UNIQUE_RECORDS_PER_SEGMENT);
for (int i = 0; i < recordRows.size(); i++) {
Object[] values = recordRows.get(i);
assertNotNull(values);
assertTrue(values.length >= 1);
int[] intValues = (int[]) values[0];
int baseValue = intValues[0];
int[] expectedConcatInt = new int[]{baseValue, baseValue + MV_OFFSET, baseValue, baseValue + MV_OFFSET};
assertTrue(Arrays.equals(intValues, expectedConcatInt));
}
}
{
// Test nesting arrayConcatInt()
String query =
getConcatNestingQueryString("arrayConcatInt", "mvIntCol", "mvRawIntCol", "concatIntsLen", "testTable",
NUM_UNIQUE_RECORDS_PER_SEGMENT);
// After concatenation, make sure each multivalue entry has double the entries
assertTrue(validateMvArrayLength(query, 2 * MV_LENGTH, NUM_UNIQUE_RECORDS_PER_SEGMENT));
}
{
// Test use of the output in WHERE expression
String query =
getConcatUseInWhereQueryString("arrayConcatInt", "mvIntCol", "mvRawIntCol", "concatInts", "testTable", "<=",
2 * MV_LENGTH, NUM_UNIQUE_RECORDS_PER_SEGMENT);
assertTrue(validateQueryOutputSize(query, NUM_UNIQUE_RECORDS_PER_SEGMENT));
query =
getConcatUseInWhereQueryString("arrayConcatInt", "mvIntCol", "mvRawIntCol", "concatInts", "testTable", ">",
2 * MV_LENGTH, NUM_UNIQUE_RECORDS_PER_SEGMENT);
assertTrue(validateQueryOutputSize(query, 0));
}
{
// Test use of output in GROUP BY
String query =
getConcatGroupByQueryString("arrayConcatInt", "mvIntCol", "mvRawIntCol", "concatInts", "testTable", "<=",
2 * MV_LENGTH, NUM_UNIQUE_RECORDS_PER_SEGMENT);
assertTrue(validateQueryOutputSize(query, NUM_UNIQUE_RECORDS_PER_SEGMENT));
}
}

@Test
public void testArrayConcatLong() {
{
// Test basic arrayConcatLong() functionality
String query = getConcatQueryString("arrayConcatLong", "mvLongCol", "mvRawLongCol", "concatLongs", "testTable",
NUM_UNIQUE_RECORDS_PER_SEGMENT);
ResultTable resultTable = getBrokerResponse(query).getResultTable();
assertNotNull(resultTable);
List<Object[]> recordRows = resultTable.getRows();
assertNotNull(recordRows);
assertEquals(recordRows.size(), NUM_UNIQUE_RECORDS_PER_SEGMENT);
for (int i = 0; i < recordRows.size(); i++) {
Object[] values = recordRows.get(i);
assertNotNull(values);
assertTrue(values.length >= 1);
long[] intValues = (long[]) values[0];
long baseValue = intValues[0];
long[] expectedConcatInt = new long[]{baseValue, baseValue + MV_OFFSET, baseValue, baseValue + MV_OFFSET};
assertTrue(Arrays.equals(intValues, expectedConcatInt));
}
}
{
// Test nesting arrayConcatLong()
String query =
getConcatNestingQueryString("arrayConcatLong", "mvLongCol", "mvRawLongCol", "concatLongsLen", "testTable",
NUM_UNIQUE_RECORDS_PER_SEGMENT);
// After concatenation, make sure each multivalue entry has double the entries
assertTrue(validateMvArrayLength(query, 2 * MV_LENGTH, NUM_UNIQUE_RECORDS_PER_SEGMENT));
}
{
// Test use of the output in WHERE expression
String query =
getConcatUseInWhereQueryString("arrayConcatLong", "mvLongCol", "mvRawLongCol", "concatLongs", "testTable",
"<=", 2 * MV_LENGTH, NUM_UNIQUE_RECORDS_PER_SEGMENT);
assertTrue(validateQueryOutputSize(query, NUM_UNIQUE_RECORDS_PER_SEGMENT));
query = getConcatUseInWhereQueryString("arrayConcatLong", "mvLongCol", "mvRawLongCol", "concatLongs", "testTable",
">", 2 * MV_LENGTH, NUM_UNIQUE_RECORDS_PER_SEGMENT);
assertTrue(validateQueryOutputSize(query, 0));
}
{
// Test use of output in GROUP BY
String query =
getConcatGroupByQueryString("arrayConcatLong", "mvLongCol", "mvRawLongCol", "concatLongs", "testTable", "<=",
2 * MV_LENGTH, NUM_UNIQUE_RECORDS_PER_SEGMENT);
assertTrue(validateQueryOutputSize(query, NUM_UNIQUE_RECORDS_PER_SEGMENT));
}
}

@Test
public void testArrayConcatFloat() {
{
// Test basic arrayConcatFloat() functionality
String query =
getConcatQueryString("arrayConcatFloat", "mvFloatCol", "mvRawFloatCol", "concatFloats", "testTable",
NUM_UNIQUE_RECORDS_PER_SEGMENT);
ResultTable resultTable = getBrokerResponse(query).getResultTable();
assertNotNull(resultTable);
List<Object[]> recordRows = resultTable.getRows();
assertNotNull(recordRows);
assertEquals(recordRows.size(), NUM_UNIQUE_RECORDS_PER_SEGMENT);
for (int i = 0; i < recordRows.size(); i++) {
Object[] values = recordRows.get(i);
assertNotNull(values);
assertTrue(values.length >= 1);
float[] intValues = (float[]) values[0];
float baseValue = intValues[0];
float[] expectedConcatInt = new float[]{baseValue, baseValue + MV_OFFSET, baseValue, baseValue + MV_OFFSET};
assertTrue(Arrays.equals(intValues, expectedConcatInt));
}
}
{
// Test nesting arrayConcatFloat()
String query =
getConcatNestingQueryString("arrayConcatFloat", "mvFloatCol", "mvRawFloatCol", "concatFloatsLen", "testTable",
NUM_UNIQUE_RECORDS_PER_SEGMENT);
// After concatenation, make sure each multivalue entry has double the entries
assertTrue(validateMvArrayLength(query, 2 * MV_LENGTH, NUM_UNIQUE_RECORDS_PER_SEGMENT));
}
{
// Test use of the output in WHERE expression
String query =
getConcatUseInWhereQueryString("arrayConcatFloat", "mvFloatCol", "mvRawFloatCol", "concatFloats", "testTable",
"<=", 2 * MV_LENGTH, NUM_UNIQUE_RECORDS_PER_SEGMENT);
assertTrue(validateQueryOutputSize(query, NUM_UNIQUE_RECORDS_PER_SEGMENT));
query =
getConcatUseInWhereQueryString("arrayConcatFloat", "mvFloatCol", "mvRawFloatCol", "concatFloats", "testTable",
">", 2 * MV_LENGTH, NUM_UNIQUE_RECORDS_PER_SEGMENT);
assertTrue(validateQueryOutputSize(query, 0));
}
{
// Test use of output in GROUP BY
String query =
getConcatGroupByQueryString("arrayConcatFloat", "mvFloatCol", "mvRawFloatCol", "concatFloats", "testTable",
"<=", 2 * MV_LENGTH, NUM_UNIQUE_RECORDS_PER_SEGMENT);
assertTrue(validateQueryOutputSize(query, NUM_UNIQUE_RECORDS_PER_SEGMENT));
}
}

@Test
public void testArrayConcatDouble() {
{
// Test basic arrayConcatDouble() functionality
String query =
getConcatQueryString("arrayConcatDouble", "mvDoubleCol", "mvRawDoubleCol", "concatDoubles", "testTable",
NUM_UNIQUE_RECORDS_PER_SEGMENT);
ResultTable resultTable = getBrokerResponse(query).getResultTable();
assertNotNull(resultTable);
List<Object[]> recordRows = resultTable.getRows();
assertNotNull(recordRows);
assertEquals(recordRows.size(), NUM_UNIQUE_RECORDS_PER_SEGMENT);
for (int i = 0; i < recordRows.size(); i++) {
Object[] values = recordRows.get(i);
assertNotNull(values);
assertTrue(values.length >= 1);
double[] intValues = (double[]) values[0];
double baseValue = intValues[0];
double[] expectedConcatInt = new double[]{baseValue, baseValue + MV_OFFSET, baseValue, baseValue + MV_OFFSET};
assertTrue(Arrays.equals(intValues, expectedConcatInt));
}
}
{
// Test nesting arrayConcatDouble()
String query =
getConcatNestingQueryString("arrayConcatDouble", "mvDoubleCol", "mvRawDoubleCol", "concatDoublesLen",
"testTable", NUM_UNIQUE_RECORDS_PER_SEGMENT);
// After concatenation, make sure each multivalue entry has double the entries
assertTrue(validateMvArrayLength(query, 2 * MV_LENGTH, NUM_UNIQUE_RECORDS_PER_SEGMENT));
}
{
// Test use of the output in WHERE expression
String query =
getConcatUseInWhereQueryString("arrayConcatDouble", "mvDoubleCol", "mvRawDoubleCol", "concatDoubles",
"testTable", "<=", 2 * MV_LENGTH, NUM_UNIQUE_RECORDS_PER_SEGMENT);
assertTrue(validateQueryOutputSize(query, NUM_UNIQUE_RECORDS_PER_SEGMENT));
query = getConcatUseInWhereQueryString("arrayConcatDouble", "mvDoubleCol", "mvRawDoubleCol", "concatDoubles",
"testTable", ">", 2 * MV_LENGTH, NUM_UNIQUE_RECORDS_PER_SEGMENT);
assertTrue(validateQueryOutputSize(query, 0));
}
{
// Test use of output in GROUP BY
String query = getConcatGroupByQueryString("arrayConcatDouble", "mvDoubleCol", "mvRawDoubleCol", "concatDoubles",
"testTable", "<=", 2 * MV_LENGTH, NUM_UNIQUE_RECORDS_PER_SEGMENT);
assertTrue(validateQueryOutputSize(query, NUM_UNIQUE_RECORDS_PER_SEGMENT));
}
}

@Test
public void testArrayConcatString() {
{
// Test basic arrayConcatDouble() functionality
String query =
getConcatQueryString("arrayConcatString", "mvStringCol", "mvRawStringCol", "concatStrings", "testTable",
NUM_UNIQUE_RECORDS_PER_SEGMENT);
ResultTable resultTable = getBrokerResponse(query).getResultTable();
assertNotNull(resultTable);
List<Object[]> recordRows = resultTable.getRows();
assertNotNull(recordRows);
assertEquals(recordRows.size(), NUM_UNIQUE_RECORDS_PER_SEGMENT);
for (int i = 0; i < recordRows.size(); i++) {
Object[] values = recordRows.get(i);
assertNotNull(values);
assertTrue(values.length >= 1);
String[] stringValues = (String[]) values[0];
assertTrue(stringValues.length == 2 * MV_LENGTH);
String[] expectedConcatStrings = new String[2 * MV_LENGTH];
for (int j = 0, k = MV_LENGTH; j < MV_LENGTH; j++) {
expectedConcatStrings[j] = stringValues[j];
expectedConcatStrings[k++] = stringValues[j];
}
assertTrue(Arrays.equals(stringValues, expectedConcatStrings));
}
}
{
// Test nesting arrayConcatString()
String query =
getConcatNestingQueryString("arrayConcatString", "mvStringCol", "mvRawStringCol", "concatStringLen",
"testTable", NUM_UNIQUE_RECORDS_PER_SEGMENT);
// After concatenation, make sure each multivalue entry has double the entries
assertTrue(validateMvArrayLength(query, 2 * MV_LENGTH, NUM_UNIQUE_RECORDS_PER_SEGMENT));
}
{
// Test use of the output in WHERE expression
String query =
getConcatUseInWhereQueryString("arrayConcatString", "mvStringCol", "mvRawStringCol", "concatStrings",
"testTable", "<=", 2 * MV_LENGTH, NUM_UNIQUE_RECORDS_PER_SEGMENT);
assertTrue(validateQueryOutputSize(query, NUM_UNIQUE_RECORDS_PER_SEGMENT));
query = getConcatUseInWhereQueryString("arrayConcatString", "mvStringCol", "mvRawStringCol", "concatStrings",
"testTable", ">", 2 * MV_LENGTH, NUM_UNIQUE_RECORDS_PER_SEGMENT);
assertTrue(validateQueryOutputSize(query, 0));
}
{
// Test use of output in GROUP BY
String query = getConcatGroupByQueryString("arrayConcatString", "mvStringCol", "mvRawStringCol", "concatStrings",
"testTable", "<=", 2 * MV_LENGTH, NUM_UNIQUE_RECORDS_PER_SEGMENT);
assertTrue(validateQueryOutputSize(query, NUM_UNIQUE_RECORDS_PER_SEGMENT));
}
}
}