Allow aliasing of Macros and add new alias for complex decode 64 (#15034)

* Add AliasExprMacro to allow aliasing of native expression macros
* Add decode_base64_complex alias for complex_decode_base64
This commit is contained in:
Pranav 2023-10-05 16:24:36 -07:00 committed by GitHub
parent 36d7b3cc65
commit 06c5527c85
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 139 additions and 35 deletions

View File

@ -251,7 +251,7 @@ public class ArrayOfDoublesSketchSqlAggregatorTest extends BaseCalciteQueryTest
+ " SUM(cnt),\n" + " SUM(cnt),\n"
+ " DS_TUPLE_DOUBLES_METRICS_SUM_ESTIMATE(DS_TUPLE_DOUBLES(tuplesketch_dim2)) AS all_sum_estimates,\n" + " DS_TUPLE_DOUBLES_METRICS_SUM_ESTIMATE(DS_TUPLE_DOUBLES(tuplesketch_dim2)) AS all_sum_estimates,\n"
+ StringUtils.replace( + StringUtils.replace(
"DS_TUPLE_DOUBLES_METRICS_SUM_ESTIMATE(DS_TUPLE_DOUBLES_INTERSECT(COMPLEX_DECODE_BASE64('arrayOfDoublesSketch', '%s'), DS_TUPLE_DOUBLES(tuplesketch_dim2), 128)) AS intersect_sum_estimates\n", "DS_TUPLE_DOUBLES_METRICS_SUM_ESTIMATE(DS_TUPLE_DOUBLES_INTERSECT(DECODE_BASE64_COMPLEX('arrayOfDoublesSketch', '%s'), DS_TUPLE_DOUBLES(tuplesketch_dim2), 128)) AS intersect_sum_estimates\n",
"%s", "%s",
COMPACT_BASE_64_ENCODED_SKETCH_FOR_INTERSECTION COMPACT_BASE_64_ENCODED_SKETCH_FOR_INTERSECTION
) )

View File

@ -32,7 +32,13 @@ public class BuiltInExprMacros
public static class ComplexDecodeBase64ExprMacro implements ExprMacroTable.ExprMacro public static class ComplexDecodeBase64ExprMacro implements ExprMacroTable.ExprMacro
{ {
public static final String NAME = "complex_decode_base64"; public static final String NAME = "complex_decode_base64";
public static final String ALIAS = "decode_base64_complex";
/**
* use name() in closure scope to allow Alias macro to override it with alias.
*
* @return String
*/
@Override @Override
public String name() public String name()
{ {
@ -52,7 +58,7 @@ public class BuiltInExprMacros
public ComplexDecodeBase64Expression(List<Expr> args) public ComplexDecodeBase64Expression(List<Expr> args)
{ {
super(NAME, args); super(name(), args);
validationHelperCheckArgumentCount(args, 2); validationHelperCheckArgumentCount(args, 2);
final Expr arg0 = args.get(0); final Expr arg0 = args.get(0);
@ -148,7 +154,6 @@ public class BuiltInExprMacros
public static class StringDecodeBase64UTFExprMacro implements ExprMacroTable.ExprMacro public static class StringDecodeBase64UTFExprMacro implements ExprMacroTable.ExprMacro
{ {
public static final String NAME = "decode_base64_utf8"; public static final String NAME = "decode_base64_utf8";
@Override @Override
@ -158,6 +163,11 @@ public class BuiltInExprMacros
return new StringDecodeBase64UTFExpression(args.get(0)); return new StringDecodeBase64UTFExpression(args.get(0));
} }
/**
* use name() in closure scope to allow Alias macro to override it with alias.
*
* @return String
*/
@Override @Override
public String name() public String name()
{ {
@ -168,7 +178,7 @@ public class BuiltInExprMacros
{ {
public StringDecodeBase64UTFExpression(Expr arg) public StringDecodeBase64UTFExpression(Expr arg)
{ {
super(NAME, arg); super(name(), arg);
} }
@Override @Override
@ -214,4 +224,5 @@ public class BuiltInExprMacros
} }
} }
} }
} }

View File

@ -43,8 +43,13 @@ import java.util.stream.Collectors;
*/ */
public class ExprMacroTable public class ExprMacroTable
{ {
private static final BuiltInExprMacros.ComplexDecodeBase64ExprMacro COMPLEX_DECODE_BASE_64_EXPR_MACRO = new BuiltInExprMacros.ComplexDecodeBase64ExprMacro();
private static final List<ExprMacro> BUILT_IN = ImmutableList.of( private static final List<ExprMacro> BUILT_IN = ImmutableList.of(
new BuiltInExprMacros.ComplexDecodeBase64ExprMacro(), COMPLEX_DECODE_BASE_64_EXPR_MACRO,
new AliasExprMacro(
COMPLEX_DECODE_BASE_64_EXPR_MACRO,
BuiltInExprMacros.ComplexDecodeBase64ExprMacro.ALIAS
),
new BuiltInExprMacros.StringDecodeBase64UTFExprMacro() new BuiltInExprMacros.StringDecodeBase64UTFExprMacro()
); );
private static final ExprMacroTable NIL = new ExprMacroTable(Collections.emptyList()); private static final ExprMacroTable NIL = new ExprMacroTable(Collections.emptyList());
@ -247,4 +252,32 @@ public class ExprMacroTable
return StringUtils.format("(%s %s)", name, getArgs()); return StringUtils.format("(%s %s)", name, getArgs());
} }
} }
/***
* Alias Expression macro to create an alias and delegate operations to the same base macro.
* The Expr spit out by the apply method should use name() in all the places instead of an internal constant so that things like error messages behave correctly.
*/
static class AliasExprMacro implements ExprMacroTable.ExprMacro
{
private final ExprMacroTable.ExprMacro exprMacro;
private final String alias;
public AliasExprMacro(final ExprMacroTable.ExprMacro baseExprMacro, final String alias)
{
this.exprMacro = baseExprMacro;
this.alias = alias;
}
@Override
public Expr apply(List<Expr> args)
{
return exprMacro.apply(args);
}
@Override
public String name()
{
return alias;
}
}
} }

View File

@ -41,6 +41,9 @@ import javax.annotation.Nullable;
import java.math.BigDecimal; import java.math.BigDecimal;
import java.math.RoundingMode; import java.math.RoundingMode;
import java.nio.ByteBuffer; import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Set; import java.util.Set;
public class FunctionTest extends InitializedNullHandlingTest public class FunctionTest extends InitializedNullHandlingTest
@ -952,8 +955,47 @@ public class FunctionTest extends InitializedNullHandlingTest
), ),
expected expected
); );
// test with alias
assertExpr(
StringUtils.format(
"decode_base64_complex('%s', '%s')",
TypeStrategiesTest.NULLABLE_TEST_PAIR_TYPE.getComplexTypeName(),
StringUtils.encodeBase64String(bytes)
),
expected
);
} }
@Test
public void testMacrosWithMultipleAliases()
{
ExprMacroTable.ExprMacro drinkMacro = new ExprMacroTable.ExprMacro()
{
@Override
public Expr apply(List<Expr> args)
{
return new StringExpr("happiness");
}
@Override
public String name()
{
return "drink";
}
};
List<ExprMacroTable.ExprMacro> macros = new ArrayList<>();
macros.add(drinkMacro);
List<String> aliases = Arrays.asList("tea", "coffee", "chai", "chaha", "kevha", "chay");
for (String tea : aliases) {
macros.add(new ExprMacroTable.AliasExprMacro(drinkMacro, tea));
}
final ExprMacroTable exprMacroTable = new ExprMacroTable(macros);
final Expr happiness = new StringExpr("happiness");
Assert.assertEquals(happiness, Parser.parse("drink(1,2)", exprMacroTable));
for (String tea : aliases) {
Assert.assertEquals(happiness, Parser.parse(StringUtils.format("%s(1,2)", tea), exprMacroTable));
}
}
@Test @Test
public void testComplexDecodeNull() public void testComplexDecodeNull()
{ {
@ -964,6 +1006,13 @@ public class FunctionTest extends InitializedNullHandlingTest
), ),
null null
); );
assertExpr(
StringUtils.format(
"decode_base64_complex('%s', null)",
TypeStrategiesTest.NULLABLE_TEST_PAIR_TYPE.getComplexTypeName()
),
null
);
} }
@Test @Test

View File

@ -32,6 +32,7 @@ import org.apache.calcite.sql.fun.SqlStdOperatorTable;
import org.apache.calcite.sql.validate.SqlNameMatcher; import org.apache.calcite.sql.validate.SqlNameMatcher;
import org.apache.druid.java.util.common.ISE; import org.apache.druid.java.util.common.ISE;
import org.apache.druid.java.util.common.StringUtils; import org.apache.druid.java.util.common.StringUtils;
import org.apache.druid.math.expr.BuiltInExprMacros;
import org.apache.druid.sql.calcite.aggregation.SqlAggregator; import org.apache.druid.sql.calcite.aggregation.SqlAggregator;
import org.apache.druid.sql.calcite.aggregation.builtin.ArrayConcatSqlAggregator; import org.apache.druid.sql.calcite.aggregation.builtin.ArrayConcatSqlAggregator;
import org.apache.druid.sql.calcite.aggregation.builtin.ArraySqlAggregator; import org.apache.druid.sql.calcite.aggregation.builtin.ArraySqlAggregator;
@ -227,11 +228,16 @@ public class DruidOperatorTable implements SqlOperatorTable
.add(ContainsOperatorConversion.caseInsensitive()) .add(ContainsOperatorConversion.caseInsensitive())
.build(); .build();
private static final SqlOperatorConversion COMPLEX_DECODE_OPERATOR_CONVERSIONS = new ComplexDecodeBase64OperatorConversion();
private static final List<SqlOperatorConversion> VALUE_COERCION_OPERATOR_CONVERSIONS = private static final List<SqlOperatorConversion> VALUE_COERCION_OPERATOR_CONVERSIONS =
ImmutableList.<SqlOperatorConversion>builder() ImmutableList.<SqlOperatorConversion>builder()
.add(new CastOperatorConversion()) .add(new CastOperatorConversion())
.add(new ReinterpretOperatorConversion()) .add(new ReinterpretOperatorConversion())
.add(new ComplexDecodeBase64OperatorConversion()) .add(COMPLEX_DECODE_OPERATOR_CONVERSIONS)
.add(new AliasedOperatorConversion(
COMPLEX_DECODE_OPERATOR_CONVERSIONS,
BuiltInExprMacros.ComplexDecodeBase64ExprMacro.ALIAS
))
.add(new DecodeBase64UTFOperatorConversion()) .add(new DecodeBase64UTFOperatorConversion())
.build(); .build();

View File

@ -30,6 +30,7 @@ import org.apache.druid.java.util.common.DateTimes;
import org.apache.druid.java.util.common.HumanReadableBytes; import org.apache.druid.java.util.common.HumanReadableBytes;
import org.apache.druid.java.util.common.Intervals; import org.apache.druid.java.util.common.Intervals;
import org.apache.druid.java.util.common.JodaUtils; import org.apache.druid.java.util.common.JodaUtils;
import org.apache.druid.java.util.common.StringUtils;
import org.apache.druid.java.util.common.UOE; import org.apache.druid.java.util.common.UOE;
import org.apache.druid.java.util.common.granularity.Granularities; import org.apache.druid.java.util.common.granularity.Granularities;
import org.apache.druid.java.util.common.granularity.PeriodGranularity; import org.apache.druid.java.util.common.granularity.PeriodGranularity;
@ -14424,36 +14425,40 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
public void testComplexDecode() public void testComplexDecode()
{ {
cannotVectorize(); cannotVectorize();
testQuery( for (String complexDecode : Arrays.asList("COMPLEX_DECODE_BASE64", "DECODE_BASE64_COMPLEX")) {
"SELECT COMPLEX_DECODE_BASE64('hyperUnique',PARSE_JSON(TO_JSON_STRING(unique_dim1))) from druid.foo LIMIT 10", testQuery(
ImmutableList.of( StringUtils.format(
Druids.newScanQueryBuilder() "SELECT %s('hyperUnique',PARSE_JSON(TO_JSON_STRING(unique_dim1))) from druid.foo LIMIT 10",
.dataSource(CalciteTests.DATASOURCE1) complexDecode
.intervals(querySegmentSpec(Filtration.eternity())) ),
.columns("v0") ImmutableList.of(
.virtualColumns( Druids.newScanQueryBuilder()
expressionVirtualColumn( .dataSource(CalciteTests.DATASOURCE1)
"v0", .intervals(querySegmentSpec(Filtration.eternity()))
"complex_decode_base64('hyperUnique',parse_json(to_json_string(\"unique_dim1\")))", .columns("v0")
ColumnType.ofComplex("hyperUnique") .virtualColumns(
) expressionVirtualColumn(
) "v0",
.resultFormat(ResultFormat.RESULT_FORMAT_COMPACTED_LIST) "complex_decode_base64('hyperUnique',parse_json(to_json_string(\"unique_dim1\")))",
.legacy(false) ColumnType.ofComplex("hyperUnique")
.limit(10) )
.build() )
), .resultFormat(ResultFormat.RESULT_FORMAT_COMPACTED_LIST)
ImmutableList.of( .legacy(false)
new Object[]{"\"AQAAAEAAAA==\""}, .limit(10)
new Object[]{"\"AQAAAQAAAAHNBA==\""}, .build()
new Object[]{"\"AQAAAQAAAAOzAg==\""}, ),
new Object[]{"\"AQAAAQAAAAFREA==\""}, ImmutableList.of(
new Object[]{"\"AQAAAQAAAACyEA==\""}, new Object[]{"\"AQAAAEAAAA==\""},
new Object[]{"\"AQAAAQAAAAEkAQ==\""} new Object[]{"\"AQAAAQAAAAHNBA==\""},
) new Object[]{"\"AQAAAQAAAAOzAg==\""},
); new Object[]{"\"AQAAAQAAAAFREA==\""},
new Object[]{"\"AQAAAQAAAACyEA==\""},
new Object[]{"\"AQAAAQAAAAEkAQ==\""}
)
);
}
} }
@Test @Test
public void testComplexDecodeAgg() public void testComplexDecodeAgg()
{ {