WIP

apache · Jan 7, 2024 · 6b65419 · 6b65419
1 parent 25c9069
commit 6b65419
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 75 deletions.
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
@@ -18,8 +18,6 @@
 package org.apache.drill.exec.expr.fn.impl.conv;
 
 
-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;
@@ -33,6 +31,8 @@
 import org.apache.drill.exec.server.options.OptionManager;
 import org.apache.drill.exec.vector.complex.writer.BaseWriter;
 import org.apache.drill.exec.vector.complex.writer.BaseWriter.ComplexWriter;
+
+import javax.inject.Inject;
 @SuppressWarnings("unused")
 public class JsonConvertFrom {
 
@@ -55,14 +55,14 @@ public static class ConvertFromJsonNullableInput implements DrillSimpleFunc {
     ResultSetLoader rsLoader;
 
     @Workspace
-    org.apache.drill.exec.store.easy.json.loader.SingleElementIterator<java.io.InputStream> streamIter;
+    org.apache.drill.exec.store.easy.json.loader.ClosingStreamIterator streamIter;
 
     @Workspace
     org.apache.drill.exec.store.easy.json.loader.JsonLoaderImpl jsonLoader;
 
     @Override
     public void setup() {
-      streamIter = new org.apache.drill.exec.store.easy.json.loader.SingleElementIterator<>();
+      streamIter = new org.apache.drill.exec.store.easy.json.loader.ClosingStreamIterator();
       rsLoader.startBatch();
     }
 
@@ -109,7 +109,7 @@ public static class ConvertFromJsonVarcharInput implements DrillSimpleFunc {
     ComplexWriter writer;
 
     @Workspace
-    org.apache.drill.exec.store.easy.json.loader.SingleElementIterator<java.io.InputStream> streamIter;
+    org.apache.drill.exec.store.easy.json.loader.ClosingStreamIterator streamIter;
 
     @Inject
     OptionManager options;
@@ -122,7 +122,7 @@ public static class ConvertFromJsonVarcharInput implements DrillSimpleFunc {
 
     @Override
     public void setup() {
-      streamIter = new org.apache.drill.exec.store.easy.json.loader.SingleElementIterator<>();
+      streamIter = new org.apache.drill.exec.store.easy.json.loader.ClosingStreamIterator();
       rsLoader.startBatch();
     }
 

diff --git a/exec/java-exec/src/main/java/org/apache/drill/exec/expr/fn/impl/conv/JsonConverterUtils.java b/exec/java-exec/src/main/java/org/apache/drill/exec/expr/fn/impl/conv/JsonConverterUtils.java
@@ -21,30 +21,13 @@
 
 import org.apache.drill.exec.physical.resultSet.ResultSetLoader;
 import org.apache.drill.exec.server.options.OptionManager;
+import org.apache.drill.exec.store.easy.json.loader.ClosingStreamIterator;
 import org.apache.drill.exec.store.easy.json.loader.JsonLoaderImpl;
 import org.apache.drill.exec.store.easy.json.loader.JsonLoaderImpl.JsonLoaderBuilder;
-import org.apache.drill.exec.store.easy.json.loader.SingleElementIterator;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
-import java.io.InputStream;
 
 public class JsonConverterUtils {
 
-  private static final Logger logger = LoggerFactory.getLogger(JsonConverterUtils.class);
-
-  /*
-  public static InputStream convertStringToInputStream(String input) {
-    try (InputStream stream = IOUtils.toInputStream(input, Charset.defaultCharset())) {
-      return stream;
-    } catch (IOException e) {
-      throw UserException.dataReadError(e)
-        .message("Unable to read JSON string")
-        .build(logger);
-    }
-  }
-  */
-
   /**
    * Creates a {@link JsonLoaderImpl} for use in JSON conversion UDFs.
    * @param rsLoader The {@link ResultSetLoader} used in the UDF
@@ -54,7 +37,7 @@ public static InputStream convertStringToInputStream(String input) {
    */
   public static JsonLoaderImpl createJsonLoader(ResultSetLoader rsLoader,
                                                 OptionManager options,
-                                                SingleElementIterator<InputStream> stream) {
+                                                ClosingStreamIterator stream) {
     // Add JSON configuration from Storage plugin, if present.
     JsonLoaderBuilder jsonLoaderBuilder = new JsonLoaderBuilder()
         .resultSetLoader(rsLoader)
@@ -63,5 +46,4 @@ public static JsonLoaderImpl createJsonLoader(ResultSetLoader rsLoader,
 
     return (JsonLoaderImpl) jsonLoaderBuilder.build();
   }
-
 }
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
@@ -20,14 +20,19 @@
 
 
 import ch.qos.logback.classic.Level;
+import org.apache.drill.common.types.TypeProtos.MinorType;
 import org.apache.drill.exec.ExecConstants;
 import org.apache.drill.exec.physical.impl.project.ProjectRecordBatch;
 import org.apache.drill.exec.physical.impl.validate.IteratorValidatorBatchIterator;
 import org.apache.drill.exec.physical.rowSet.RowSet;
+import org.apache.drill.exec.physical.rowSet.RowSetBuilder;
+import org.apache.drill.exec.record.metadata.SchemaBuilder;
+import org.apache.drill.exec.record.metadata.TupleMetadata;
 import org.apache.drill.exec.store.easy.json.loader.JsonLoaderImpl;
 import org.apache.drill.test.ClusterFixture;
 import org.apache.drill.test.ClusterTest;
 import org.apache.drill.test.LogFixture;
+import org.apache.drill.test.rowSet.RowSetComparison;
 import org.junit.BeforeClass;
 import org.junit.Test;
 
@@ -61,61 +66,27 @@ public void testConvertFromJsonVarBinary() throws Exception {
 
   @Test
   public void testConvertFromJsonVarChar() throws Exception {
-    // String sql = "SELECT *, convert_FromJSON('{\"foo\":\"bar\"}') FROM cp.`jsoninput/allTypes.csv`";
-    String sql = "SELECT convert_FromJSON('{\"foo\":\"bar\"}') FROM (VALUES(1))";
+    String sql = "SELECT json_data['foo'] AS foo, json_data['num'] AS num FROM " +
+        "(SELECT convert_FromJSON('{\"foo\":\"bar\", \"num\":10}') as json_data FROM (VALUES(1)))";
     RowSet results = client.queryBuilder().sql(sql).rowSet();
-    results.clear();
-  }
 
-/*
-  private void doTestConvertToJsonFunction() throws Exception {
-    String table = "nan_test.csv";
-    File file = new File(dirTestWatcher.getRootDir(), table);
-    String csv = "col_0, {\"nan_col\":NaN}";
-    String query = String.format("select string_binary(convert_toJSON(convert_fromJSON(columns[1]))) as col " +
-      "from dfs.`%s` where columns[0]='col_0'", table);
-    try {
-      FileUtils.writeStringToFile(file, csv,  Charset.defaultCharset());
-      List<QueryDataBatch> results = testSqlWithResults(query);
-      RecordBatchLoader batchLoader = new RecordBatchLoader(getAllocator());
-      assertEquals("Query result must contain 1 row", 1, results.size());
-      QueryDataBatch batch = results.get(0);
+    TupleMetadata expectedSchema = new SchemaBuilder()
+        .addNullable("foo", MinorType.VARCHAR)
+        .addNullable("num", MinorType.BIGINT)
+        .buildSchema();
+
+    RowSet expected = new RowSetBuilder(client.allocator(), expectedSchema)
+        .addRow("bar", 10L)
+        .build();
 
-      batchLoader.load(batch.getHeader().getDef(), batch.getData());
-      VectorWrapper<?> vw = batchLoader.getValueAccessorById(VarCharVector.class, batchLoader.getValueVectorId(SchemaPath.getCompoundPath("col")).getFieldIds());
-      // ensuring that `NaN` token ARE NOT enclosed with double quotes
-      String resultJson = vw.getValueVector().getAccessor().getObject(0).toString();
-      int nanIndex = resultJson.indexOf("NaN");
-      assertNotEquals("`NaN` must not be enclosed with \"\" ", '"', resultJson.charAt(nanIndex - 1));
-      assertNotEquals("`NaN` must not be enclosed with \"\" ", '"', resultJson.charAt(nanIndex + "NaN".length()));
-      batch.release();
-      batchLoader.clear();
-    } finally {
-      FileUtils.deleteQuietly(file);
-    }
+    new RowSetComparison(expected).verifyAndClearAll(results);
   }
 
   @Test
-  public void testConvertFromJsonFunction() throws Exception {
-    //runBoth(this::doTestConvertFromJsonFunction);
-  }
+  public void testMultipleRows() throws Exception {
+    String sql = "SELECT string_binary(convert_toJSON(`name`)) FROM cp.`jsoninput/multirow.csvh`";
 
-  private void doTestConvertFromJsonFunction() throws Exception {
-    String table = "nan_test.csv";
-    File file = new File(dirTestWatcher.getRootDir(), table);
-    String csv = "col_0, {\"nan_col\":NaN}";
-    try {
-      FileUtils.writeStringToFile(file, csv);
-      testBuilder()
-        .sqlQuery(String.format("select convert_fromJSON(columns[1]) as col from dfs.`%s`", table))
-        .unOrdered()
-        .baselineColumns("col")
-        .baselineValues(mapOf("nan_col", Double.NaN))
-        .go();
-    } finally {
-      FileUtils.deleteQuietly(file);
-    }
+    RowSet results = client.queryBuilder().sql(sql).rowSet();
+    results.print();
   }
-  */
-
 }
diff --git a/exec/java-exec/src/test/resources/jsoninput/multirow.csvh b/exec/java-exec/src/test/resources/jsoninput/multirow.csvh
@@ -0,0 +1,3 @@
+num, name
+1, "bob"
+4,"steve"