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]