Skip to content

Commit

Permalink
Added additional optional args for the function
Browse files Browse the repository at this point in the history
  • Loading branch information
@cgivre
cgivre committed Jul 7, 2024
1 parent 7ddde7b commit 153ef52
Show file tree
Hide file tree
Showing 2 changed files with 163 additions and 0 deletions.
101 changes: 101 additions & 0 deletions exec/java-exec/src/main/java/org/apache/drill/exec/expr/fn/impl/conv/JsonConvertFrom.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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 {

Expand Down Expand Up @@ -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);
}
}
}

}
62 changes: 62 additions & 0 deletions exec/java-exec/src/test/java/org/apache/drill/exec/store/json/TestJsonConversionUDF.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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);
}
}

0 comments on commit 153ef52

Please sign in to comment.