New jsonl parser implementation and benchmark

Author: asgerf

extensions/ql-vscode/src/common/jsonl-reader.ts

@@ -1,4 +1,4 @@
-import { statSync } from "fs";
+import { stat } from "fs/promises";
 import { createReadStream } from "fs-extra";
 import { createInterface } from "readline";
 
@@ -33,7 +33,7 @@ export async function readJsonlFile<T>(
   }
 
   void logger?.log(
-    `Parsing ${path} (${statSync(path).size / 1024 / 1024} MB)...`,
+    `Parsing ${path} (${(await stat(path)).size / 1024 / 1024} MB)...`,
   );
   const fileStream = createReadStream(path, "utf8");
   const rl = createInterface({
@@ -64,3 +64,48 @@ export async function readJsonlFile<T>(
   }
   logProgress();
 }
+
+const doubleLineBreakRegexp = /\n\r?\n/;
+
+export async function readJsonlFile2<T>(
+  path: string,
+  handler: (value: T) => Promise<void>,
+  logger?: { log: (message: string) => void },
+): Promise<void> {
+  void logger?.log(
+    `Parsing ${path} (${(await stat(path)).size / 1024 / 1024} MB)...`,
+  );
+  return new Promise((resolve, reject) => {
+    const stream = createReadStream(path, { encoding: "utf8" });
+    let buffer = "";
+    stream.on("data", async (chunk: string) => {
+      const parts = (buffer + chunk).split(doubleLineBreakRegexp);
+      buffer = parts.pop()!;
+      if (parts.length > 0) {
+        try {
+          stream.pause();
+          for (const part of parts) {
+            await handler(JSON.parse(part));
+          }
+          stream.resume();
+        } catch (e) {
+          stream.destroy();
+          reject(e);
+        }
+      }
+    });
+    stream.on("end", async () => {
+      if (buffer.trim().length > 0) {
+        try {
+          await handler(JSON.parse(buffer));
+        } catch (e) {
+          reject(e);
+          return;
+        }
+      }
+      void logger?.log(`Finishing parsing ${path}`);
+      resolve();
+    });
+    stream.on("error", reject);
+  });
+}

extensions/ql-vscode/test/benchmarks/jsonl-reader.bench.ts

@@ -0,0 +1,74 @@
+import { readFile } from "fs-extra";
+import { readJsonlFile, readJsonlFile2 } from "../../src/common/jsonl-reader";
+import { performance } from "perf_hooks";
+import { join } from "path";
+
+/** An "obviously correct" implementation to test against. */
+async function readJsonlReferenceImpl<T>(
+  path: string,
+  handler: (value: T) => Promise<void>,
+): Promise<void> {
+  const logSummary = await readFile(path, "utf-8");
+
+  // Remove newline delimiters because summary is in .jsonl format.
+  const jsonSummaryObjects: string[] = logSummary.split(/\r?\n\r?\n/g);
+
+  for (const obj of jsonSummaryObjects) {
+    const jsonObj = JSON.parse(obj) as T;
+    await handler(jsonObj);
+  }
+}
+
+type ParserFn = (
+  text: string,
+  callback: (v: unknown) => Promise<void>,
+) => Promise<void>;
+
+const parsers: Record<string, ParserFn> = {
+  readJsonlReferenceImpl,
+  readJsonlFile,
+  readJsonlFile2,
+};
+
+async function main() {
+  const args = process.argv.slice(2);
+  const file =
+    args.length > 0
+      ? args[0]
+      : join(
+          __dirname,
+          "../unit-tests/data/evaluator-log-summaries/bad-join-order.jsonl",
+        );
+  const numTrials = args.length > 1 ? Number(args[1]) : 100;
+  const referenceValues: any[] = [];
+  await readJsonlReferenceImpl(file, async (event) => {
+    referenceValues.push(event);
+  });
+  const referenceValueString = JSON.stringify(referenceValues);
+  // Do warm-up runs and check against reference implementation
+  for (const [name, parser] of Object.entries(parsers)) {
+    const values: unknown[] = [];
+    await parser(file, async (event) => {
+      values.push(event);
+    });
+    if (JSON.stringify(values) !== referenceValueString) {
+      console.error(`${name}: failed to match reference implementation`);
+    }
+  }
+  for (const [name, parser] of Object.entries(parsers)) {
+    const startTime = performance.now();
+    for (let i = 0; i < numTrials; ++i) {
+      await Promise.all([
+        parser(file, async () => {}),
+        parser(file, async () => {}),
+      ]);
+    }
+    const duration = performance.now() - startTime;
+    const durationPerTrial = duration / numTrials;
+    console.log(`${name}: ${durationPerTrial.toFixed(1)} ms`);
+  }
+}
+
+main().catch((err: unknown) => {
+  console.error(err);
+});