Merge pull request ZJONSSON#1 from arnabguptadev/fixBrokenWrite

Fix for bug that causes memory overrun and adds more rows than it should
thenickdude · Jun 10, 2020 · 0c2d106 · 0c2d106
2 parents 641cdbc + caec575
commit 0c2d106
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 6 deletions.
diff --git a/lib/writer.js b/lib/writer.js
@@ -94,8 +94,11 @@ class ParquetWriter {
     parquet_shredder.shredRecord(this.schema, row, this.rowBuffer);
 
     if (this.rowBuffer.rowCount >= this.rowGroupSize) {
-      await this.envelopeWriter.writeRowGroup(this.rowBuffer);
+      let to_write = this.rowBuffer;
+      let rowCount = this.rowBuffer.rowCount;
       this.rowBuffer = {};
+      to_write.rowCount = rowCount;
+      await this.envelopeWriter.writeRowGroup(to_write);
     }
   }
 
@@ -112,9 +115,12 @@ class ParquetWriter {
 
     this.closed = true;
 
-    if (this.rowBuffer.rowCount > 0 || this.rowBuffer.rowCount >= this.rowGroupSize) {
-      await this.envelopeWriter.writeRowGroup(this.rowBuffer);
+    if (this.rowBuffer.rowCount > 0) {
+      let to_write = this.rowBuffer;
+      let rowCount = this.rowBuffer.rowCount;
       this.rowBuffer = {};
+      to_write.rowCount = rowCount;
+      await this.envelopeWriter.writeRowGroup(to_write);
     }
 
     await this.envelopeWriter.writeFooter(this.userMetadata);

diff --git a/test/integration.js b/test/integration.js
@@ -105,9 +105,7 @@ async function writeTestFile(opts) {
 
   let rows = mkTestRows(opts);
 
-  for (let row of rows) {
-    await writer.appendRow(row);
-  }
+  rows.forEach(async row => await writer.appendRow(row));
 
   await writer.close();
 }