LakshSingla commented on code in PR #16572:
URL: https://github.com/apache/druid/pull/16572#discussion_r1642222664


##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/results/ExportResultsFrameProcessor.java:
##########
@@ -167,9 +175,22 @@ private void exportFrame(final Frame frame)
               for (int j = 0; j < exportRowSignature.size(); j++) {
                 String columnName = exportRowSignature.getColumnName(j);
                 BaseObjectColumnValueSelector<?> selector = 
selectors.get(outputColumnNameToFrameColumnNumberMap.getInt(columnName));
-                exportWriter.writeRowField(columnName, selector.getObject());
+                if (resultsContext == null) {
+                  exportWriter.writeRowField(columnName, selector.getObject());

Review Comment:
   Maybe it should fail if resultsContext == null to avoid correctness issues, 
rather than being backward compatible? 



##########
extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQExportTest.java:
##########
@@ -133,7 +133,142 @@ public void testNumberOfRowsPerFile()
 
     Assert.assertEquals(
         expectedFooFileContents(false).size() + 1, // + 1 for the manifest file
-        Objects.requireNonNull(new 
File(exportDir.getAbsolutePath()).listFiles()).length
+        Objects.requireNonNull(exportDir.listFiles()).length
+    );
+  }
+
+  @Test
+  void testExportComplexColumns() throws IOException
+  {
+    final RowSignature rowSignature = RowSignature.builder()
+                                                  .add("__time", 
ColumnType.LONG)
+                                                  .add("a", ColumnType.LONG)
+                                                  .add("b", ColumnType.LONG)
+                                                  .add("c_json", 
ColumnType.STRING).build();
+
+    final File exportDir = newTempFolder("export");
+    final String sql = StringUtils.format("INSERT INTO\n"
+                                          + "EXTERN(local(exportPath=>'%s'))\n"
+                                          + "AS CSV\n"
+                                          + "SELECT\n"
+                                          + "  \"a\",\n"
+                                          + "  \"b\",\n"
+                                          + "  json_object(key 'c' value b) 
c_json\n"
+                                          + "FROM (\n"
+                                          + "  SELECT *\n"
+                                          + "  FROM TABLE(\n"
+                                          + "    EXTERN(\n"
+                                          + "      
'{\"type\":\"inline\",\"data\":\"a,b\\n1,1\\n2,2\"}',\n"
+                                          + "      
'{\"type\":\"csv\",\"findColumnsFromHeader\":true}'\n"
+                                          + "    )\n"
+                                          + "  ) EXTEND (\"a\" BIGINT, \"b\" 
BIGINT)\n"
+                                          + ")", exportDir.getAbsolutePath());
+
+    testIngestQuery().setSql(sql)
+                     .setExpectedDataSource("foo1")
+                     .setQueryContext(DEFAULT_MSQ_CONTEXT)
+                     .setExpectedRowSignature(rowSignature)
+                     .setExpectedSegment(ImmutableSet.of())
+                     .setExpectedResultRows(ImmutableList.of())
+                     .verifyResults();
+
+    Assert.assertEquals(
+        2, // result file and manifest file
+        Objects.requireNonNull(exportDir.listFiles()).length
+    );
+
+    File resultFile = new File(exportDir, 
"query-test-query-worker0-partition0.csv");
+    List<String> results = readResultsFromFile(resultFile);
+    Assert.assertEquals(
+        ImmutableList.of(
+            "a,b,c_json", "1,1,\"{\"\"c\"\":1}\"", "2,2,\"{\"\"c\"\":2}\""
+        ),
+        results
+    );
+  }
+
+  @Test
+  void testExportSketchColumns() throws IOException
+  {
+    final RowSignature rowSignature = RowSignature.builder()
+                                                  .add("__time", 
ColumnType.LONG)
+                                                  .add("a", ColumnType.LONG)
+                                                  .add("b", ColumnType.LONG)
+                                                  .add("c_json", 
ColumnType.STRING).build();
+
+    final File exportDir = newTempFolder("export");
+    final String sql = StringUtils.format("INSERT INTO\n"
+                                          + "EXTERN(local(exportPath=>'%s'))\n"
+                                          + "AS CSV\n"
+                                          + "SELECT\n"
+                                          + "  \"a\",\n"
+                                          + "  \"b\",\n"
+                                          + "  ds_hll(b) c_ds_hll\n"
+                                          + "FROM (\n"
+                                          + "  SELECT *\n"
+                                          + "  FROM TABLE(\n"
+                                          + "    EXTERN(\n"
+                                          + "      
'{\"type\":\"inline\",\"data\":\"a,b\\n1,b1\\n2,b2\"}',\n"
+                                          + "      
'{\"type\":\"csv\",\"findColumnsFromHeader\":true}'\n"
+                                          + "    )\n"
+                                          + "  ) EXTEND (\"a\" BIGINT, \"b\" 
VARCHAR)\n"
+                                          + ")\n"
+                                          + "GROUP BY 1,2", 
exportDir.getAbsolutePath());
+
+    testIngestQuery().setSql(sql)
+                     .setExpectedDataSource("foo1")
+                     .setQueryContext(DEFAULT_MSQ_CONTEXT)
+                     .setExpectedRowSignature(rowSignature)
+                     .setExpectedSegment(ImmutableSet.of())
+                     .setExpectedResultRows(ImmutableList.of())
+                     .verifyResults();
+
+    Assert.assertEquals(
+        2, // result file and manifest file
+        Objects.requireNonNull(exportDir.listFiles()).length
+    );
+
+    File resultFile = new File(exportDir, 
"query-test-query-worker0-partition0.csv");
+    List<String> results = readResultsFromFile(resultFile);
+    Assert.assertEquals(
+        ImmutableList.of(
+            "a,b,c_ds_hll", "1,b1,\"\"\"AgEHDAMIAQBa1y0L\"\"\"", 
"2,b2,\"\"\"AgEHDAMIAQCi6V0G\"\"\""
+        ),
+        results
+    );
+  }
+
+  @Test
+  void testEmptyExport() throws IOException
+  {
+    RowSignature rowSignature = RowSignature.builder()
+                                            .add("__time", ColumnType.LONG)
+                                            .add("dim1", ColumnType.STRING)
+                                            .add("cnt", 
ColumnType.LONG).build();
+
+    File exportDir = newTempFolder("export");
+    final String sql = StringUtils.format("insert into 
extern(local(exportPath=>'%s')) as csv select cnt, dim1 as dim from foo where 
dim1='nonexistentvalue'", exportDir.getAbsolutePath());

Review Comment:
   nit: SQL capitalization and spacing isn't same as the one added in the 
`testExportSketchColumns`. 



##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/results/ExportResultsFrameProcessorFactory.java:
##########
@@ -63,19 +64,22 @@ public class ExportResultsFrameProcessorFactory implements 
FrameProcessorFactory
   private final ExportStorageProvider exportStorageProvider;
   private final ResultFormat exportFormat;
   private final ColumnMappings columnMappings;
+  private final ResultsContext resultsContext;
 
   @JsonCreator
   public ExportResultsFrameProcessorFactory(
       @JsonProperty("queryId") String queryId,
       @JsonProperty("exportStorageProvider") ExportStorageProvider 
exportStorageProvider,
       @JsonProperty("exportFormat") ResultFormat exportFormat,
-      @JsonProperty("columnMappings") @Nullable ColumnMappings columnMappings
+      @JsonProperty("columnMappings") @Nullable ColumnMappings columnMappings,
+      @JsonProperty("resultsContext") @Nullable ResultsContext resultsContext

Review Comment:
   Are we maintaining backward compatibility for EXPORT? This change doesn't 
look like it's backward compatible. 



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to