[GitHub] vvysotskyi commented on a change in pull request #1549: DRILL-6834: Introduce option to disable result set on CTAS, create vi…

[GitBox <gi...@apache.org>]

vvysotskyi commented on a change in pull request #1549: DRILL-6834: Introduce option to disable result set on CTAS, create vi…
URL: https://github.com/apache/drill/pull/1549#discussion_r235493106
 
 

 ##########
 File path: exec/java-exec/src/main/java/org/apache/drill/exec/physical/impl/materialize/VectorRecordMaterializer.java
 ##########
 @@ -25,35 +27,44 @@
 import org.apache.drill.exec.record.BatchSchema;
 import org.apache.drill.exec.record.RecordBatch;
 import org.apache.drill.exec.record.WritableBatch;
+import org.apache.drill.exec.server.options.OptionManager;
 
-public class VectorRecordMaterializer implements RecordMaterializer{
+public class VectorRecordMaterializer implements RecordMaterializer {
   static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(VectorRecordMaterializer.class);
 
   private QueryId queryId;
   private RecordBatch batch;
   private BufferAllocator allocator;
+  private OptionManager options;
 
   public VectorRecordMaterializer(FragmentContext context, OperatorContext oContext, RecordBatch batch) {
     this.queryId = context.getHandle().getQueryId();
     this.batch = batch;
     this.allocator = oContext.getAllocator();
     BatchSchema schema = batch.getSchema();
     assert schema != null : "Schema must be defined.";
-
-//    for (MaterializedField f : batch.getSchema()) {
-//      logger.debug("New Field: {}", f);
-//    }
+    options = context.getOptions();
   }
 
   public QueryWritableBatch convertNext() {
-    //batch.getWritableBatch().getDef().getRecordCount()
     WritableBatch w = batch.getWritableBatch().transfer(allocator);
+    QueryData.Builder builder = QueryData.newBuilder()
+        .setQueryId(queryId)
+        .setRowCount(batch.getRecordCount())
+        .setDef(w.getDef());
+    if (hasUpdateCount()) {
+      int count = w.getDef().getUpdateCount();
+      builder.setUpdateCount(count == -1 ? 0 : count);
+    }
+    return new QueryWritableBatch(builder.build(), w.getBuffers());
+  }
 
-    QueryData header = QueryData.newBuilder() //
-        .setQueryId(queryId) //
-        .setRowCount(batch.getRecordCount()) //
-        .setDef(w.getDef()).build();
-    QueryWritableBatch batch = new QueryWritableBatch(header, w.getBuffers());
-    return batch;
+  private boolean hasUpdateCount() {
 
 Review comment:
   Could you please rename this method to something more descriptive.

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services