From 153ef526474a3e199516daf6f7cadc5f96b05c60 Mon Sep 17 00:00:00 2001 From: cgivre Date: Sun, 7 Jul 2024 11:49:50 -0400 Subject: [PATCH] Added additional optional args for the function --- .../expr/fn/impl/conv/JsonConvertFrom.java | 101 ++++++++++++++++++ .../store/json/TestJsonConversionUDF.java | 62 +++++++++++ 2 files changed, 163 insertions(+) diff --git a/exec/java-exec/src/main/java/org/apache/drill/exec/expr/fn/impl/conv/JsonConvertFrom.java b/exec/java-exec/src/main/java/org/apache/drill/exec/expr/fn/impl/conv/JsonConvertFrom.java index 7b7a82ed3bc..ac27db7951b 100644 --- a/exec/java-exec/src/main/java/org/apache/drill/exec/expr/fn/impl/conv/JsonConvertFrom.java +++ b/exec/java-exec/src/main/java/org/apache/drill/exec/expr/fn/impl/conv/JsonConvertFrom.java @@ -25,6 +25,7 @@ import org.apache.drill.exec.expr.annotations.Output; import org.apache.drill.exec.expr.annotations.Param; import org.apache.drill.exec.expr.annotations.Workspace; +import org.apache.drill.exec.expr.holders.BitHolder; import org.apache.drill.exec.expr.holders.NullableVarBinaryHolder; import org.apache.drill.exec.expr.holders.NullableVarCharHolder; import org.apache.drill.exec.expr.holders.VarBinaryHolder; @@ -123,6 +124,53 @@ public void eval() { } } + @FunctionTemplate(name = "convert_fromJSON", scope = FunctionScope.SIMPLE, isRandom = true) + public static class ConvertFromJsonVarcharWithConfig implements DrillSimpleFunc { + + @Param + VarCharHolder in; + + @Param + BitHolder allTextModeHolder; + + @Param + BitHolder readNumbersAsDoubleHolder; + + @Inject + DrillBuf buffer; + + @Workspace + org.apache.drill.exec.vector.complex.fn.JsonReader jsonReader; + + @Output + ComplexWriter writer; + + @Override + public void setup() { + boolean allTextMode = allTextModeHolder.value == 1; + boolean readNumbersAsDouble = readNumbersAsDoubleHolder.value == 1; + + jsonReader = new org.apache.drill.exec.vector.complex.fn.JsonReader.Builder(buffer) + .defaultSchemaPathColumns() + .allTextMode(allTextMode) + .readNumbersAsDouble(readNumbersAsDouble) + .build(); + } + + @Override + public void eval() { + try { + jsonReader.setSource(in.start, in.end, in.buffer); + jsonReader.write(writer); + buffer = jsonReader.getWorkBuf(); + } catch (Exception e) { + throw new org.apache.drill.common.exceptions.DrillRuntimeException("Error while converting from JSON. ", e); + } + } + } + + + @FunctionTemplate(name = "convert_fromJSON", scope = FunctionScope.SIMPLE, isRandom = true) public static class ConvertFromJsonNullableInput implements DrillSimpleFunc { @@ -221,4 +269,57 @@ public void eval() { } } } + + @FunctionTemplate(name = "convert_fromJSON", scope = FunctionScope.SIMPLE, isRandom = true) + public static class ConvertFromJsonVarcharNullableInputWithConfigs implements DrillSimpleFunc { + + @Param + NullableVarCharHolder in; + + @Param + BitHolder allTextModeHolder; + + @Param + BitHolder readNumbersAsDoubleHolder; + + @Inject + DrillBuf buffer; + + @Workspace + org.apache.drill.exec.vector.complex.fn.JsonReader jsonReader; + + @Output ComplexWriter writer; + + @Override + public void setup() { + boolean allTextMode = allTextModeHolder.value == 1; + boolean readNumbersAsDouble = readNumbersAsDoubleHolder.value == 1; + + jsonReader = new org.apache.drill.exec.vector.complex.fn.JsonReader.Builder(buffer) + .defaultSchemaPathColumns() + .allTextMode(allTextMode) + .readNumbersAsDouble(readNumbersAsDouble) + .build(); + } + + @Override + public void eval() { + if (in.isSet == 0) { + // Return empty map + org.apache.drill.exec.vector.complex.writer.BaseWriter.MapWriter mapWriter = writer.rootAsMap(); + mapWriter.start(); + mapWriter.end(); + return; + } + + try { + jsonReader.setSource(in.start, in.end, in.buffer); + jsonReader.write(writer); + buffer = jsonReader.getWorkBuf(); + } catch (Exception e) { + throw new org.apache.drill.common.exceptions.DrillRuntimeException("Error while converting from JSON. ", e); + } + } + } + } diff --git a/exec/java-exec/src/test/java/org/apache/drill/exec/store/json/TestJsonConversionUDF.java b/exec/java-exec/src/test/java/org/apache/drill/exec/store/json/TestJsonConversionUDF.java index 32c918eda43..d6fa857a390 100644 --- a/exec/java-exec/src/test/java/org/apache/drill/exec/store/json/TestJsonConversionUDF.java +++ b/exec/java-exec/src/test/java/org/apache/drill/exec/store/json/TestJsonConversionUDF.java @@ -122,6 +122,38 @@ public void testAllTextMode() throws Exception { resetSessionOption(ExecConstants.JSON_READ_NUMBERS_AS_DOUBLE); } + @Test + public void testAllTextModeFromArgs() throws Exception { + // Set the system options to make sure that the UDF is using the provided options rather than + // the system options. + alterSession(ExecConstants.JSON_ALL_TEXT_MODE, false); + alterSession(ExecConstants.JSON_READ_NUMBERS_AS_DOUBLE, true); + + String sql = "SELECT \n" + + "typeof(jsonMap['bi']) AS bi, \n" + + "typeof(jsonMap['fl']) AS fl, \n" + + "typeof(jsonMap['st']) AS st, \n" + + "typeof(jsonMap['mp']) AS mp, \n" + + "typeof(jsonMap['ar']) AS ar, \n" + + "typeof(jsonMap['nu']) AS nu\n" + + "FROM(\n" + + "SELECT convert_fromJSON(col1, true, false) AS jsonMap, \n" + + "col2 \n" + + "FROM cp.`jsoninput/allTypes.csvh`\n" + + ")"; + + testBuilder() + .sqlQuery(sql) + .unOrdered() + .baselineColumns("bi", "fl", "st", "mp", "ar", "nu") + .baselineValues("VARCHAR", "VARCHAR", "VARCHAR", "MAP", "VARCHAR", "NULL") + .go(); + + resetSessionOption(ExecConstants.JSON_ALL_TEXT_MODE); + resetSessionOption(ExecConstants.JSON_READ_NUMBERS_AS_DOUBLE); + } + + @Test public void testReadNumbersAsDouble() throws Exception { alterSession(ExecConstants.JSON_ALL_TEXT_MODE, false); @@ -158,4 +190,34 @@ public void testReadNumbersAsDouble() throws Exception { resetSessionOption(ExecConstants.JSON_ALL_TEXT_MODE); resetSessionOption(ExecConstants.JSON_READ_NUMBERS_AS_DOUBLE); } + + @Test + public void testReadNumbersAsDoubleFromArgs() throws Exception { + // Set the system options to make sure that the UDF is using the provided options rather than + // the system options. + alterSession(ExecConstants.JSON_ALL_TEXT_MODE, true); + alterSession(ExecConstants.JSON_READ_NUMBERS_AS_DOUBLE, true); + String sql = "SELECT \n" + + "typeof(jsonMap['bi']) AS bi, \n" + + "typeof(jsonMap['fl']) AS fl, \n" + + "typeof(jsonMap['st']) AS st, \n" + + "typeof(jsonMap['mp']) AS mp, \n" + + "typeof(jsonMap['ar']) AS ar, \n" + + "typeof(jsonMap['nu']) AS nu\n" + + "FROM(\n" + + "SELECT convert_fromJSON(col1, false, true) AS jsonMap, \n" + + "col2 \n" + + "FROM cp.`jsoninput/allTypes.csvh`\n" + + ")"; + + testBuilder() + .sqlQuery(sql) + .unOrdered() + .baselineColumns("bi", "fl", "st", "mp", "ar", "nu") + .baselineValues("FLOAT8", "FLOAT8", "VARCHAR", "MAP", "FLOAT8", "NULL") + .go(); + + resetSessionOption(ExecConstants.JSON_ALL_TEXT_MODE); + resetSessionOption(ExecConstants.JSON_READ_NUMBERS_AS_DOUBLE); + } }