Skip to content

Commit

Permalink
DRILL-8501: Json Conversion UDF Not Respecting System JSON Options (#…
Browse files Browse the repository at this point in the history 
…2921)
  • Loading branch information
@cgivre
cgivre committed Jul 9, 2024
1 parent af62ace commit d44aa98
Show file tree
Hide file tree
Showing 4 changed files with 516 additions and 101 deletions.
291 changes: 271 additions & 20 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 @@ -19,41 +19,99 @@


import io.netty.buffer.DrillBuf;

import javax.inject.Inject;

import org.apache.drill.exec.expr.DrillSimpleFunc;
import org.apache.drill.exec.expr.annotations.FunctionTemplate;
import org.apache.drill.exec.expr.annotations.FunctionTemplate.FunctionScope;
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;
import org.apache.drill.exec.expr.holders.VarCharHolder;
import org.apache.drill.exec.server.options.OptionManager;
import org.apache.drill.exec.vector.complex.writer.BaseWriter.ComplexWriter;

public class JsonConvertFrom {
import javax.inject.Inject;

static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(JsonConvertFrom.class);
public class JsonConvertFrom {

private JsonConvertFrom() {
}

@FunctionTemplate(name = "convert_fromJSON", scope = FunctionScope.SIMPLE, isRandom = true)
public static class ConvertFromJson implements DrillSimpleFunc {

@Param VarBinaryHolder in;
@Inject DrillBuf buffer;
@Workspace org.apache.drill.exec.vector.complex.fn.JsonReader jsonReader;
@Param
VarBinaryHolder in;

@Output ComplexWriter writer;
@Inject
DrillBuf buffer;

@Inject
OptionManager options;

@Workspace
org.apache.drill.exec.vector.complex.fn.JsonReader jsonReader;

@Output
ComplexWriter writer;

@Override
public void setup() {
boolean allTextMode = options.getBoolean(org.apache.drill.exec.ExecConstants.JSON_ALL_TEXT_MODE);
boolean readNumbersAsDouble = options.getBoolean(org.apache.drill.exec.ExecConstants.JSON_READ_NUMBERS_AS_DOUBLE);

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 ConvertFromJsonWithArgs implements DrillSimpleFunc {

@Param
VarBinaryHolder 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();
}

Expand All @@ -69,19 +127,79 @@ public void eval() {
}
}


@FunctionTemplate(name = "convert_fromJSON", scope = FunctionScope.SIMPLE, isRandom = true)
public static class ConvertFromJsonVarchar implements DrillSimpleFunc {

@Param VarCharHolder in;
@Inject DrillBuf buffer;
@Workspace org.apache.drill.exec.vector.complex.fn.JsonReader jsonReader;
@Param
VarCharHolder in;

@Output ComplexWriter writer;
@Inject
DrillBuf buffer;

@Inject
OptionManager options;

@Workspace
org.apache.drill.exec.vector.complex.fn.JsonReader jsonReader;

@Output
ComplexWriter writer;

@Override
public void setup() {
boolean allTextMode = options.getBoolean(org.apache.drill.exec.ExecConstants.JSON_ALL_TEXT_MODE);
boolean readNumbersAsDouble = options.getBoolean(org.apache.drill.exec.ExecConstants.JSON_READ_NUMBERS_AS_DOUBLE);

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 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();
}

Expand All @@ -100,16 +218,83 @@ public void eval() {
@FunctionTemplate(name = "convert_fromJSON", scope = FunctionScope.SIMPLE, isRandom = true)
public static class ConvertFromJsonNullableInput implements DrillSimpleFunc {

@Param NullableVarBinaryHolder in;
@Inject DrillBuf buffer;
@Workspace org.apache.drill.exec.vector.complex.fn.JsonReader jsonReader;
@Param
NullableVarBinaryHolder in;

@Output ComplexWriter writer;
@Inject
DrillBuf buffer;

@Inject
OptionManager options;

@Workspace
org.apache.drill.exec.vector.complex.fn.JsonReader jsonReader;

@Output
ComplexWriter writer;

@Override
public void setup() {
boolean allTextMode = options.getBoolean(org.apache.drill.exec.ExecConstants.JSON_ALL_TEXT_MODE);
boolean readNumbersAsDouble = options.getBoolean(org.apache.drill.exec.ExecConstants.JSON_READ_NUMBERS_AS_DOUBLE);

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);
}
}
}

@FunctionTemplate(name = "convert_fromJSON", scope = FunctionScope.SIMPLE, isRandom = true)
public static class ConvertFromJsonNullableInputWithArgs implements DrillSimpleFunc {

@Param
NullableVarBinaryHolder 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();
}

Expand All @@ -136,16 +321,29 @@ public void eval() {
@FunctionTemplate(name = "convert_fromJSON", scope = FunctionScope.SIMPLE, isRandom = true)
public static class ConvertFromJsonVarcharNullableInput implements DrillSimpleFunc {

@Param NullableVarCharHolder in;
@Inject DrillBuf buffer;
@Workspace org.apache.drill.exec.vector.complex.fn.JsonReader jsonReader;
@Param
NullableVarCharHolder in;

@Inject
DrillBuf buffer;

@Inject
OptionManager options;

@Workspace
org.apache.drill.exec.vector.complex.fn.JsonReader jsonReader;

@Output ComplexWriter writer;

@Override
public void setup() {
boolean allTextMode = options.getBoolean(org.apache.drill.exec.ExecConstants.JSON_ALL_TEXT_MODE);
boolean readNumbersAsDouble = options.getBoolean(org.apache.drill.exec.ExecConstants.JSON_READ_NUMBERS_AS_DOUBLE);

jsonReader = new org.apache.drill.exec.vector.complex.fn.JsonReader.Builder(buffer)
.defaultSchemaPathColumns()
.allTextMode(allTextMode)
.readNumbersAsDouble(readNumbersAsDouble)
.build();
}

Expand All @@ -168,4 +366,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);
}
}
}

}
Loading

0 comments on commit d44aa98

Please sign in to comment.