feat: dbm partition (#57)

* add: partition filter

* update: test

* update: interface

* fix: build

* update: interface

* bump version

Author: vpbs2

benchmarking/benchmarking-app/src/app/dbm-context/raw-dbm-context.tsx

@@ -34,9 +34,9 @@ export const RawDBMProvider = ({ children }: { children: JSX.Element }) => {
       },
     });
     /**
-     * Making the queryWithTableNames simply run the queries without sequence which is the default behavior
+     * Making the queryWithTables simply run the queries without sequence which is the default behavior
      */
-    dbm.queryWithTableNames = async ({ query, tableNames }) => {
+    dbm.queryWithTables = async ({ query, tables }) => {
       return dbm.query(query);
     };
     setdbm(dbm);

benchmarking/benchmarking-app/src/app/query-benchmarking/query-benchmarking.tsx

@@ -23,7 +23,7 @@ export const QueryBenchmarking = () => {
       const eachQueryStart = performance.now();
 
       const promiseObj = dbm
-        .queryWithTableNames({ query: TEST_QUERIES[i], tableNames: ['taxi'] })
+        .queryWithTables({ query: TEST_QUERIES[i], tables: [{ name: 'taxi' }] })
         .then((results) => {
           const end = performance.now();
           const time = end - eachQueryStart;

meerkat-dbm/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@devrev/meerkat-dbm",
-  "version": "0.0.143",
+  "version": "0.0.15",
   "dependencies": {
     "tslib": "^2.3.0",
     "@duckdb/duckdb-wasm": "^1.28.0",

meerkat-dbm/src/dbm/dbm.spec.ts

@@ -8,7 +8,7 @@ import {
 import { FileData, Table, TableWiseFiles } from '../types';
 import { DBM } from './dbm';
 import { InstanceManagerType } from './instance-manager';
-import { DBMConstructorOptions } from './types';
+import { DBMConstructorOptions, TableConfig } from './types';
 
 export class MockFileManager implements FileManagerType {
   private fileBufferStore: Record<string, FileBufferStore> = {};
@@ -53,7 +53,8 @@ export class MockFileManager implements FileManagerType {
     return fileBuffer.buffer;
   }
 
-  async mountFileBufferByTableNames(tableNames: string[]): Promise<void> {
+  async mountFileBufferByTables(tables: TableConfig[]): Promise<void> {
+    const tableNames = tables.map((table) => table.name);
     for (const tableName of tableNames) {
       for (const key in this.fileBufferStore) {
         if (this.fileBufferStore[key].tableName === tableName) {
@@ -64,17 +65,6 @@ export class MockFileManager implements FileManagerType {
     }
   }
 
-  async unmountFileBufferByTableNames(tableNames: string[]): Promise<void> {
-    for (const tableName of tableNames) {
-      for (const key in this.fileBufferStore) {
-        if (this.fileBufferStore[key].tableName === tableName) {
-          // unmount operation here
-          console.log(`Unmounted file buffer for ${key}`);
-        }
-      }
-    }
-  }
-
   async getFilesByTableName(tableName: string): Promise<FileData[]> {
     const files: FileData[] = [];
 
@@ -96,15 +86,17 @@ export class MockFileManager implements FileManagerType {
     }
   }
 
-  async getFilesNameForTables(tableNames: string[]): Promise<TableWiseFiles[]> {
+  async getFilesNameForTables(
+    tableNames: TableConfig[]
+  ): Promise<TableWiseFiles[]> {
     const data: TableWiseFiles[] = [];
 
-    for (const tableName of tableNames) {
-      const files: string[] = [];
+    for (const { name: tableName } of tableNames) {
+      const files: FileData[] = [];
 
       for (const key in this.fileBufferStore) {
         if (this.fileBufferStore[key].tableName === tableName) {
-          files.push(key);
+          files.push({ fileName: key });
         }
       }
 
@@ -170,6 +162,8 @@ describe('DBM', () => {
   let dbm: DBM;
   let instanceManager: InstanceManager;
 
+  const tables = [{ name: 'table1' }];
+
   beforeAll(async () => {
     // eslint-disable-next-line @typescript-eslint/ban-ts-comment
     //@ts-ignore
@@ -201,7 +195,7 @@ describe('DBM', () => {
     });
   });
 
-  describe('queryWithTableNames', () => {
+  describe('queryWithTables', () => {
     it('should call the preQuery hook', async () => {
       const preQuery = jest.fn();
 
@@ -211,9 +205,9 @@ describe('DBM', () => {
         buffer: new Uint8Array(),
       });
 
-      const result = await dbm.queryWithTableNames({
+      const result = await dbm.queryWithTables({
         query: 'SELECT * FROM table1',
-        tableNames: ['table1'],
+        tables: tables,
         options: {
           preQuery,
         },
@@ -224,27 +218,27 @@ describe('DBM', () => {
       expect(preQuery).toBeCalledWith([
         {
           tableName: 'table1',
-          files: ['file1'],
+          files: [{ fileName: 'file1' }],
         },
       ]);
     });
 
     it('should execute a query with table names', async () => {
-      const result = await dbm.queryWithTableNames({
+      const result = await dbm.queryWithTables({
         query: 'SELECT * FROM table1',
-        tableNames: ['table1'],
+        tables: tables,
       });
       expect(result).toEqual(['SELECT * FROM table1']);
     });
 
     it('should execute multiple queries with table names', async () => {
-      const promise1 = dbm.queryWithTableNames({
+      const promise1 = dbm.queryWithTables({
         query: 'SELECT * FROM table1',
-        tableNames: ['table1'],
+        tables: tables,
       });
-      const promise2 = dbm.queryWithTableNames({
+      const promise2 = dbm.queryWithTables({
         query: 'SELECT * FROM table2',
-        tableNames: ['table1'],
+        tables: tables,
       });
       /**
        * Number of queries in the queue should be 1 as the first query is running
@@ -271,9 +265,9 @@ describe('DBM', () => {
       /**
        * Execute another query
        */
-      const promise3 = dbm.queryWithTableNames({
+      const promise3 = dbm.queryWithTables({
         query: 'SELECT * FROM table3',
-        tableNames: ['table1'],
+        tables: tables,
       });
 
       /**
@@ -314,17 +308,17 @@ describe('DBM', () => {
       /**
        * Execute a query
        */
-      const promise1 = dbm.queryWithTableNames({
+      const promise1 = dbm.queryWithTables({
         query: 'SELECT * FROM table1',
-        tableNames: ['table1'],
+        tables: tables,
       });
 
       /**
        * Execute another query
        */
-      const promise2 = dbm.queryWithTableNames({
+      const promise2 = dbm.queryWithTables({
         query: 'SELECT * FROM table2',
-        tableNames: ['table1'],
+        tables: tables,
       });
 
       /**
@@ -375,17 +369,17 @@ describe('DBM', () => {
       /**
        * Execute a query
        */
-      const promise1 = dbm.queryWithTableNames({
+      const promise1 = dbm.queryWithTables({
         query: 'SELECT * FROM table1',
-        tableNames: ['table1'],
+        tables: tables,
       });
 
       /**
        * Execute another query
        */
-      const promise2 = dbm.queryWithTableNames({
+      const promise2 = dbm.queryWithTables({
         query: 'SELECT * FROM table2',
-        tableNames: ['table1'],
+        tables: tables,
       });
 
       /**
@@ -420,9 +414,9 @@ describe('DBM', () => {
       /**
        * Execute a query
        */
-      await dbm.queryWithTableNames({
+      await dbm.queryWithTables({
         query: 'SELECT * FROM table1',
-        tableNames: ['table1'],
+        tables: tables,
       });
 
       /**
@@ -443,9 +437,9 @@ describe('DBM', () => {
 
       // check the current query throws error abort is emitted
       try {
-        const promise = dbm.queryWithTableNames({
+        const promise = dbm.queryWithTables({
           query: 'SELECT * FROM table1',
-          tableNames: ['table1'],
+          tables: tables,
           options: {
             signal: abortController1.signal,
           },
@@ -467,17 +461,17 @@ describe('DBM', () => {
 
       const mockDBMQuery = jest.spyOn(dbm, 'query');
 
-      const promise1 = dbm.queryWithTableNames({
+      const promise1 = dbm.queryWithTables({
         query: 'SELECT * FROM table1',
-        tableNames: ['table1'],
+        tables: tables,
         options: {
           signal: abortController1.signal,
         },
       });
 
-      const promise2 = dbm.queryWithTableNames({
+      const promise2 = dbm.queryWithTables({
         query: 'SELECT * FROM table2',
-        tableNames: ['table2'],
+        tables: [{ name: 'table2' }],
         options: {
           signal: abortController2.signal,
         },

meerkat-dbm/src/dbm/dbm.ts

@@ -4,7 +4,12 @@ import { v4 as uuidv4 } from 'uuid';
 import { FileManagerType } from '../file-manager/file-manager-type';
 import { DBMEvent, DBMLogger } from '../logger';
 import { InstanceManagerType } from './instance-manager';
-import { DBMConstructorOptions, QueryOptions, QueryQueueItem } from './types';
+import {
+  DBMConstructorOptions,
+  QueryOptions,
+  QueryQueueItem,
+  TableConfig,
+} from './types';
 
 export class DBM {
   private fileManager: FileManagerType;
@@ -94,16 +99,16 @@ export class DBM {
     return this.connection;
   }
 
-  private async _queryWithTableNames(
+  private async _queryWithTables(
     query: string,
-    tableNames: string[],
+    tables: TableConfig[],
     options?: QueryOptions
   ) {
     /**
      * Load all the files into the database
      */
     const startMountTime = Date.now();
-    await this.fileManager.mountFileBufferByTableNames(tableNames);
+    await this.fileManager.mountFileBufferByTables(tables);
     const endMountTime = Date.now();
 
     this.logger.debug(
@@ -119,9 +124,7 @@ export class DBM {
       metadata: options?.metadata,
     });
 
-    const tablesFileData = await this.fileManager.getFilesNameForTables(
-      tableNames
-    );
+    const tablesFileData = await this.fileManager.getFilesNameForTables(tables);
 
     /**
      * Execute the preQuery hook
@@ -186,6 +189,7 @@ export class DBM {
      * Get the first query
      */
     this.currentQueryItem = this.queriesQueue.shift();
+
     /**
      * If there is no query, stop the queue
      */
@@ -212,9 +216,9 @@ export class DBM {
       /**
        * Execute the query
        */
-      const result = await this._queryWithTableNames(
+      const result = await this._queryWithTables(
         this.currentQueryItem.query,
-        this.currentQueryItem.tableNames,
+        this.currentQueryItem.tables,
         this.currentQueryItem.options
       );
       const endTime = Date.now();
@@ -305,13 +309,13 @@ export class DBM {
     signal?.addEventListener('abort', signalHandler);
   }
 
-  public async queryWithTableNames({
+  public async queryWithTables({
     query,
-    tableNames,
+    tables,
     options,
   }: {
     query: string;
-    tableNames: string[];
+    tables: TableConfig[];
     options?: QueryOptions;
   }) {
     const connectionId = uuidv4();
@@ -321,7 +325,7 @@ export class DBM {
 
       this.queriesQueue.push({
         query,
-        tableNames,
+        tables,
         promise: {
           resolve,
           reject,

meerkat-dbm/src/dbm/types.ts

@@ -74,9 +74,22 @@ export interface QueryOptions {
   signal?: AbortSignal;
 }
 
+export interface TableConfig {
+  /**
+   * @description
+   * Name of the table.
+   */
+  name: string;
+  /**
+   * @description
+   * Partitions of the table.
+   */
+  partitions?: string[];
+}
+
 export interface QueryQueueItem {
   query: string;
-  tableNames: string[];
+  tables: TableConfig[];
   promise: {
     resolve: (value: any) => void;
     reject: (reason?: any) => void;