Log number of rows in a chunk (#1104)

For debugging purposes it would be useful to understand the number of
rows fetched in a chunk. I don't believe this is currently captured
anywhere in the existing logs.

Co-authored-by: Thomas Powell <tpowell@palantir.com>
Co-authored-by: Samikshya Chand <148681192+samikshya-db@users.noreply.github.com>

Author: 3 people

src/main/java/com/databricks/jdbc/api/impl/arrow/AbstractArrowResultChunk.java

@@ -65,10 +65,12 @@ public abstract class AbstractArrowResultChunk {
   static final class ArrowData {
     private final List<List<ValueVector>> valueVectors;
     private final List<String> metadata;
+    private final long rowCount;
 
-    public ArrowData(List<List<ValueVector>> valueVectors, List<String> metadata) {
+    public ArrowData(List<List<ValueVector>> valueVectors, List<String> metadata, long rowCount) {
       this.valueVectors = valueVectors;
       this.metadata = metadata;
+      this.rowCount = rowCount;
     }
 
     public List<List<ValueVector>> getValueVectors() {
@@ -78,6 +80,10 @@ public List<List<ValueVector>> getValueVectors() {
     public List<String> getMetadata() {
       return metadata;
     }
+
+    public long getRowCount() {
+      return rowCount;
+    }
   }
 
   protected AbstractArrowResultChunk(
@@ -291,7 +297,11 @@ protected void initializeData(InputStream inputStream)
     ArrowData arrowData = getRecordBatchList(inputStream, rootAllocator, statementId, chunkIndex);
     recordBatchList = arrowData.getValueVectors();
     arrowMetadata = arrowData.getMetadata();
-    LOGGER.debug("Data parsed for chunk index {} and statement {}", chunkIndex, statementId);
+    LOGGER.debug(
+        "Data parsed for chunk index {} and statement {}. Row count: {}",
+        chunkIndex,
+        statementId,
+        arrowData.getRowCount());
     setStatus(ChunkStatus.PROCESSING_SUCCEEDED);
   }
 
@@ -311,10 +321,12 @@ private ArrowData getRecordBatchList(
       throws IOException {
     List<List<ValueVector>> recordBatchList = new ArrayList<>();
     List<String> metadata = new ArrayList<>();
+    long rowCount = 0L;
     try (ArrowStreamReader arrowStreamReader = new ArrowStreamReader(inputStream, rootAllocator)) {
       VectorSchemaRoot vectorSchemaRoot = arrowStreamReader.getVectorSchemaRoot();
       boolean fetchedMetadata = false;
       while (arrowStreamReader.loadNextBatch()) {
+        rowCount += vectorSchemaRoot.getRowCount();
         if (!fetchedMetadata) {
           metadata = getMetadataInformationFromSchemaRoot(vectorSchemaRoot);
           fetchedMetadata = true;
@@ -338,7 +350,7 @@ private ArrowData getRecordBatchList(
       throw e;
     }
 
-    return new ArrowData(recordBatchList, metadata);
+    return new ArrowData(recordBatchList, metadata, rowCount);
   }
 
   private List<String> getMetadataInformationFromSchemaRoot(VectorSchemaRoot vectorSchemaRoot) {