feat(filter): add basic filtering function to streams

Author: OliverDudgeon

src/filter.ts

@@ -0,0 +1,37 @@
+import type { SDFRecord } from "./parser";
+
+export type FilterFn = (record: SDFRecord) => boolean;
+
+export type FilterRule = {
+  property: string;
+  min: number;
+  max: number;
+  treatAs: "number" | "string" | "date";
+};
+
+export const filterRecord = (record: SDFRecord, rules: FilterRule[]): boolean => {
+  for (const rule of rules) {
+    const value = record.properties[rule.property];
+    // skip a record that is missing a property to which we apply a filter
+    if (value === undefined) continue;
+
+    // handle each type of filter
+    switch (rule.treatAs) {
+      case "number": {
+        if (value === "") return false;
+        // parseFloat has issues so should probably use a different library
+        const numberValue = Number.parseFloat(value);
+        // if the value is not parsable to a number, drop the record
+        if (Number.isNaN(numberValue)) return false;
+        // Keep records within the range inclusive of values equal to the bounds, I.e [min, max]
+        if (numberValue < rule.min || numberValue > rule.max) return false;
+        break;
+      }
+      // TODO: To implement
+      case "string":
+      case "date":
+        continue;
+    }
+  }
+  return true;
+};

src/index.ts

@@ -1,3 +1,4 @@
-export * from "./node-stream";
+export * from "./filter";
+export * from "./node-stream";
 export * from "./parser";
 export * from "./web-stream";

src/node-stream.ts

@@ -1,5 +1,7 @@
 import { Transform, type TransformCallback, type TransformOptions } from "node:stream";
 
+import type { FilterFn } from "./filter";
+import type { SDFRecord } from "./parser";
 import { parseSdPart } from "./parser";
 import { splitLines } from "./utils";
 
@@ -41,16 +43,18 @@ const countRecords = (buffer: string) => buffer.match(/\${4}.*/g)?.length ?? 0;
  */
 export class NodeSDFTransformer extends Transform {
   constructor(
+    private filter: FilterFn = () => true,
     options?: TransformOptions,
+    // these shouldn't be in the constructor definition but how set these to this without ts complaining?
     private buffer = "",
+    private record: SDFRecord | undefined = undefined,
   ) {
     super({ ...options, readableObjectMode: true, writableObjectMode: true });
-    this.buffer = buffer;
 
     this.push("[");
   }
 
-  private parse() {
+  private parse(): SDFRecord {
     const recordEndIndex = this.buffer.indexOf(RECORD_SEPARATOR);
     const recordText = this.buffer.slice(0, recordEndIndex);
 
@@ -66,20 +70,30 @@ export class NodeSDFTransformer extends Transform {
 
     this.buffer += data.replace(/\r\n/g, "\n");
 
-    while (countRecords(this.buffer) > 1) {
+    while (countRecords(this.buffer) > 0) {
       const record = this.parse();
-      const json = JSON.stringify(record);
-      this.push(json + ",");
+      if (this.filter(record)) {
+        if (this.record) {
+          const json = JSON.stringify(this.record);
+          this.push(json + ",");
+          this.record = record;
+        } else {
+          this.record = record;
+        }
+      }
     }
 
     callback();
   }
 
   _flush(callback: TransformCallback) {
-    // Handle the remaining record in the buffer
-    const record = this.parse();
-    const json = JSON.stringify(record);
-    this.push(json);
+    const record = this.record;
+
+    if (record) {
+      const json = JSON.stringify(record);
+      this.push(json);
+    }
+
     this.push("]");
 
     callback();

src/parser.ts

@@ -2,7 +2,10 @@
 
 export type SDFRecord = {
   molFile: string | undefined;
-  properties: Record<string, string | undefined>;
+  // mapping from property name to property value
+  // doesn't strictly need undefined here since it's only show the property may be missing,
+  // tsconfig should
+  properties: Record<string, string>;
 };
 
 export const parseSdPart = (recordLines: string[]) => {
@@ -19,12 +22,12 @@ export const parseSdPart = (recordLines: string[]) => {
     if (record.molFile) {
       const matchHeader = line.match(patternHeader);
       if (matchHeader?.length == 2) {
-        key = matchHeader[1];
+        key = matchHeader[1] as string;
       } else {
         const hasKeyAndValue = key === "" || line === "";
         if (!hasKeyAndValue) {
           const oldValue = record.properties[key];
-          record.properties[key] = oldValue ? `${oldValue};${line}` : line;
+          record.properties[key] = oldValue ? `${oldValue}\r\n${line}` : line;
         }
       }
     } else {

src/web-stream.ts

@@ -1,3 +1,4 @@
+import type { FilterFn } from "./filter";
 import { parseSdPart, type SDFRecord } from "./parser";
 import { splitLines } from "./utils";
 
@@ -19,7 +20,7 @@ const RECORD_SEPARATOR = "$$$$";
   ```
  * @returns instance of `TransformStream`
  */
-export const createSDFTransformer = () => {
+export const createSDFTransformer = (filter: FilterFn = () => true) => {
   let content = ""; // accumulator for the content of the current record
 
   // TransformStream to be used with stream.pipeThrough()
@@ -35,7 +36,9 @@ export const createSDFTransformer = () => {
         const record = parseSdPart(recordLines);
         content = content.slice(recordEndIndex + RECORD_SEPARATOR.length + 1);
 
-        controller.enqueue(record);
+        if (filter(record)) {
+          controller.enqueue(record);
+        }
       }
     },
   });

tests/stream.spec.ts

@@ -2,8 +2,9 @@ import { Transform, Writable } from "node:stream";
 import njfetch from "node-fetch";
 import { describe, expect, it } from "vitest";
 
-import type { SDFRecord } from "..";
+import { filterRecord, type FilterRule } from "../src/filter";
 import { NodeSDFTransformer } from "../src/node-stream";
+import type { SDFRecord } from "../src/parser";
 import { createSDFTransformer as createWebSDFTransformer, parser } from "..";
 
 const webStreamParser = async (stream: ReadableStream<string>) => {
@@ -112,3 +113,154 @@ describe("NodeJS stream", async () => {
     expect(streamedRecords).toEqual(records);
   });
 });
+
+describe("Filter with treatAs number", () => {
+  const record: SDFRecord = {
+    molFile: "",
+    properties: {
+      property1: "",
+      property2: "abc",
+      property3: "123",
+      property4: "-123",
+    },
+  };
+
+  it("Filter drops empty string if treated as a number", () => {
+    const filterRules: FilterRule[] = [
+      {
+        property: "property1",
+        treatAs: "number",
+        min: 0,
+        max: Infinity,
+      },
+    ];
+    const filter = (record: SDFRecord) => filterRecord(record, filterRules);
+
+    expect(filter(record)).toBe(false);
+  });
+
+  it("Filter drops non-numbers if treated as a number", () => {
+    const filterRules: FilterRule[] = [
+      {
+        property: "property2",
+        treatAs: "number",
+        min: 0,
+        max: Infinity,
+      },
+    ];
+    const filter = (record: SDFRecord) => filterRecord(record, filterRules);
+
+    expect(filter(record)).toBe(false);
+  });
+
+  it("Filter keeps numbers when using infinite bounds", () => {
+    const filterRules: FilterRule[] = [
+      {
+        property: "property3",
+        treatAs: "number",
+        min: -Infinity,
+        max: Infinity,
+      },
+    ];
+    const filter = (record: SDFRecord) => filterRecord(record, filterRules);
+
+    expect(filter(record)).toBe(true);
+  });
+
+  it("Filter keeps numbers when using semi-infinite bounds", () => {
+    const filterRules: FilterRule[] = [
+      {
+        property: "property3",
+        treatAs: "number",
+        min: 0,
+        max: Infinity,
+      },
+    ];
+    const filter = (record: SDFRecord) => filterRecord(record, filterRules);
+
+    expect(filter(record)).toBe(true);
+  });
+
+  it("Filter drops negative numbers when out of bounds", () => {
+    const filterRules: FilterRule[] = [
+      {
+        property: "property4",
+        treatAs: "number",
+        min: 0,
+        max: Infinity,
+      },
+    ];
+    const filter = (record: SDFRecord) => filterRecord(record, filterRules);
+
+    expect(filter(record)).toBe(false);
+  });
+});
+
+const getAStream = async () => {
+  const response = await njfetch(
+    "https://github.com/InformaticsMatters/sdf-parser/raw/master/tests/data/poses.sdf",
+  );
+  const stream = response.body;
+
+  if (!stream) throw new Error("No stream");
+
+  return stream;
+};
+
+describe("NodeJS stream with filter", async () => {
+  it("Filter drops all records", async () => {
+    const stream = await getAStream();
+    const filterRules: FilterRule[] = [
+      {
+        property: "TransFSScore",
+        treatAs: "number",
+        min: -Infinity,
+        max: 0,
+      },
+    ];
+    const filter = (record: SDFRecord) => filterRecord(record, filterRules);
+
+    const streamedRecords = await consumeStream(
+      stream.pipe(decoderTransform()).pipe(new NodeSDFTransformer(filter)),
+    );
+    expect(streamedRecords).toHaveLength(0);
+  });
+
+  it("Filter keeps all records", async () => {
+    const stream = await getAStream();
+
+    const filterRules: FilterRule[] = [
+      {
+        property: "FeatureStein",
+        treatAs: "number",
+        min: -Infinity,
+        max: Infinity,
+      },
+    ];
+    const filter = (record: SDFRecord) => filterRecord(record, filterRules);
+
+    const streamedRecords = await consumeStream(
+      stream.pipe(decoderTransform()).pipe(new NodeSDFTransformer(filter)),
+    );
+    expect(streamedRecords).toHaveLength(268);
+  });
+
+  it("Filter keeps the right number of records", async () => {
+    const stream = await getAStream();
+
+    const filterRules: FilterRule[] = [
+      {
+        property: "TransFSScore",
+        treatAs: "number",
+        min: 0.2,
+        max: 0.3,
+      },
+    ];
+    const filter = (record: SDFRecord) => filterRecord(record, filterRules);
+
+    const streamedRecords = await consumeStream(
+      stream.pipe(decoderTransform()).pipe(new NodeSDFTransformer(filter)),
+    );
+    expect(streamedRecords).toHaveLength(53);
+  });
+});

tsconfig.json

@@ -11,7 +11,7 @@
     "moduleDetection": "force",
     // Strictness
     "strict": true,
-    "noUncheckedIndexedAccess": false,
+    "noUncheckedIndexedAccess": true,
     // If transpiling with TS
     // "moduleResolution": "NodeNext",
     // "module": "NodeNext",