@uwdata/mosaic-duckdb 0.1.0 → 0.2.0

package/bin/csv2arrow.js

@@ -8,7 +8,7 @@ const db = new DuckDB();
 await db.csv('data', process.argv[2]);
 
 // get output stream of arrow bytes
-const stream = await db.arrowStream('SELECT * FROM data');
+const stream = await db.arrowBuffer('SELECT * FROM data');
 
  // determine the output stream
 const output = process.argv[3]

package/bin/run-server.js

@@ -0,0 +1,3 @@
+#! /usr/bin/env node
+import { DuckDB, dataServer } from '../src/index.js';
+dataServer(new DuckDB(), { rest: true, socket: true });

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@uwdata/mosaic-duckdb",
-  "version": "0.1.0",
+  "version": "0.2.0",
   "description": "A Promise-based DuckDB API and Node.js data server.",
   "keywords": [
     "duckdb",
@@ -20,12 +20,13 @@
   },
   "scripts": {
     "lint": "eslint src test --ext .js",
+    "server": "node bin/run-server.js",
     "test": "mocha 'test/**/*-test.js'",
     "prepublishOnly": "npm run test && npm run lint"
   },
   "dependencies": {
-    "duckdb": "~0.6.1",
-    "ws": "^8.12.0"
+    "duckdb": "~0.7.1",
+    "ws": "^8.13.0"
   },
-  "gitHead": "a7967c35349bdf7f00abb113ce1dd9abb233cd62"
+  "gitHead": "e53cd914c807f99aabe78dcbe618dd9543e2f438"
 }

package/src/Cache.js

@@ -0,0 +1,113 @@
+import { createHash } from 'node:crypto';
+import fs from 'node:fs/promises';
+import path from 'node:path';
+
+const DEFAULT_CACHE_DIR = '.cache';
+const DEFAULT_TTL = 1000 * 60 * 60 * 24 * 7; // 7 days
+
+export function cacheKey(hashable, type) {
+  return createHash('sha256').update(hashable).digest('hex') + '.' + type;
+}
+
+class CacheEntry {
+  constructor(data, ttl = DEFAULT_TTL) {
+    this.data = data;
+    this.touch(ttl);
+  }
+  touch(ttl = DEFAULT_TTL) {
+    this.last = Math.round(Math.max(this.last, performance.now() + ttl));
+    return this;
+  }
+}
+
+export class Cache {
+  constructor({
+    max = 10000, // max entries
+    dir = DEFAULT_CACHE_DIR,
+    ttl = DEFAULT_TTL
+  }) {
+    this.cache = new Map;
+    this.max = max;
+    this.dir = dir;
+    this.ttl = ttl;
+    readEntries(dir, this.cache);
+  }
+
+  has(key) {
+    return this.cache.has(key);
+  }
+
+  delete(key) {
+    const deleted = this.cache.delete(key);
+    if (deleted) {
+      fs.rm(path.resolve(this.dir, key), { force: true });
+    }
+    return deleted;
+  }
+
+  get(key) {
+    return this.cache.get(key)?.touch(this.ttl).data;
+  }
+
+  set(key, data, { persist = false, ttl = this.ttl } = {}) {
+    const entry = new CacheEntry(data, persist ? Infinity : ttl);
+    this.cache.set(key, entry);
+    if (persist) writeEntry(this.dir, key, entry);
+    if (this.shouldEvict()) setTimeout(() => this.evict());
+    return this;
+  }
+
+  shouldEvict() {
+    return this.cache.size > this.max;
+  }
+
+  evict() {
+    const expire = performance.now();
+    let lruKey = null;
+    let lruLast = Infinity;
+
+    for (const [key, entry] of this.cache) {
+      const { last } = entry;
+      if (last === Infinity) continue;
+
+      // least recently used entry seen so far
+      if (last < lruLast) {
+        lruKey = key;
+        lruLast = last;
+      }
+
+      // remove if time since last access exceeds ttl
+      if (expire > last) {
+        this.cache.delete(key);
+      }
+    }
+
+    // remove lru entry
+    if (this.cache.size > this.max && lruKey) {
+      this.cache.delete(lruKey);
+    }
+  }
+}
+
+async function readEntries(dir, cache) {
+  let files;
+  try {
+    files = await fs.readdir(dir);
+  } catch (err) {
+    return; // dir does not exist, nothing to do
+  }
+  await Promise.allSettled(files.map(async file => {
+    const m = file.match(/.*\.(arrow|json)/);
+    const key = m?.[1] || null;
+    if (key) {
+      const data = await fs.readFile(path.resolve(dir, file));
+      cache.set(key, new CacheEntry(data, Infinity));
+    }
+  }));
+}
+
+function writeEntry(dir, key, entry) {
+  return fs.mkdir(dir, { recursive: true }).then(
+    () => fs.writeFile(path.resolve(dir, key), entry.data)
+  );
+}

package/src/DuckDB.js

@@ -1,9 +1,10 @@
 import duckdb from 'duckdb';
-import { readFile } from 'node:fs/promises';
 import { mergeBuffers } from './merge-buffers.js';
 
+const TEMP_DIR = '.duckdb';
+
 const CONFIG = [
-  `PRAGMA temp_directory='./duckdb.tmp'`,
+  `PRAGMA temp_directory='${TEMP_DIR}'`,
   `INSTALL arrow`,
   `INSTALL httpfs`,
   `LOAD arrow`,
@@ -29,35 +30,6 @@ export class DuckDB {
     });
   }
 
-  async csv(tableName, fileName, options = {}) {
-    const opt = Object.entries({ sample_size: -1, ...options })
-      .map(([key, value]) => {
-        const t = typeof value;
-        const v = t === 'boolean' ? String(value).toUpperCase()
-          : t === 'string' ? `'${value}'`
-          : value;
-        return `${key.toUpperCase()}=${v}`;
-      })
-      .join(', ');
-    return this.exec(`CREATE TABLE ${tableName} AS SELECT *
-      FROM read_csv_auto('${fileName}', ${opt});`);
-  }
-
-  async parquet(tableName, fileName) {
-    return this.exec(`CREATE TABLE ${tableName} AS SELECT *
-      FROM read_parquet('${fileName}');`);
-  }
-
-  async ipc(tableName, buffer) {
-    const bufName = `__ipc__${tableName}`;
-    const arrowData = ArrayBuffer.isView(buffer) ? buffer : await readFile(buffer);
-    this.con.register_buffer(bufName, [arrowData], true, err => {
-      if (err) console.error(err);
-    });
-    await this.exec(`CREATE TABLE ${tableName} AS SELECT * FROM ${bufName}`);
-    this.con.unregister_buffer(bufName);
-  }
-
   prepare(sql) {
     return new DuckDBStatement(this.con.prepare(sql));
   }
@@ -97,10 +69,6 @@ export class DuckDB {
       });
     });
   }
-
-  arrowStream(sql) {
-    return this.con.arrowIPCStream(sql);
-  }
 }
 
 export class DuckDBStatement {
@@ -151,8 +119,4 @@ export class DuckDBStatement {
       });
     });
   }
-
-  arrowStream(params) {
-    return this.statement.arrowIPCStream(...params);
-  }
 }

package/src/data-server.js

@@ -1,13 +1,21 @@
 import http from 'node:http';
+import path from 'node:path';
 import url from 'node:url';
 import { WebSocketServer } from 'ws';
+import { Cache, cacheKey } from './Cache.js';
+import { createBundle, loadBundle } from './load/bundle.js';
+
+const CACHE_DIR = '.mosaic/cache';
+const BUNDLE_DIR = '.mosaic/bundle';
 
 export function dataServer(db, {
+  cache = true,
   rest = true,
   socket = true,
   port = 3000
 } = {}) {
-  const handleQuery = queryHandler(db);
+  const queryCache = cache ? new Cache({ dir: CACHE_DIR }) : null;
+  const handleQuery = queryHandler(db, queryCache);
   const app = createHTTPServer(handleQuery, rest);
   if (socket) createSocketServer(app, handleQuery);
 
@@ -25,7 +33,6 @@ function createHTTPServer(handleQuery, rest) {
       return;
     }
 
-    resp.setHeader('Content-Type', 'application/json');
     resp.setHeader('Access-Control-Allow-Origin', '*');
     resp.setHeader('Access-Control-Request-Method', '*');
     resp.setHeader('Access-Control-Allow-Methods', 'OPTIONS, POST, GET');
@@ -61,7 +68,27 @@ function createSocketServer(server, handleQuery) {
   });
 }
 
-function queryHandler(db) {
+function queryHandler(db, queryCache) {
+
+  // retrieve query result
+  async function retrieve(query, get) {
+    const { sql, type, persist } = query;
+    const key = cacheKey(sql, type);
+    let result = queryCache?.get(key);
+
+    if (result) {
+      console.log('CACHE HIT');
+    } else {
+      result = await get(sql);
+      if (persist) {
+        queryCache?.set(key, result, { persist });
+      }
+    }
+
+    return result;
+  }
+
+  // query request handler
   return async (res, data) => {
     const t0 = performance.now();
 
@@ -75,64 +102,56 @@ function queryHandler(db) {
     }
 
     try {
-      const { sql, type } = query;
-      console.log('QUERY', sql);
-
-      // request the lock to serialize requests
-      // we do this to avoid DuckDB + Arrow errors
-      await res.lock?.();
+      const { sql, type = 'json' } = query;
+      console.log(`> ${type.toUpperCase()}${sql ? ' ' + sql : ''}`);
 
       // process query and return result
       switch (type) {
-        case 'arrow':
-          // Apache Arrow response format
-          await res.stream(await db.arrowStream(sql));
-          break;
         case 'exec':
           // Execute query with no return value
           await db.exec(sql);
           res.done();
           break;
-        default:
+        case 'arrow':
+          // Apache Arrow response format
+          res.arrow(await retrieve(query, sql => db.arrowBuffer(sql)));
+          break;
+        case 'json':
           // JSON response format
-          res.json(await db.query(sql));
+          res.json(await retrieve(query, sql => db.query(sql)));
+          break;
+        case 'create-bundle':
+          // Create a named bundle of precomputed resources
+          await createBundle(
+            db, queryCache, query.queries,
+            path.resolve(BUNDLE_DIR, query.name)
+          );
+          res.done();
+          break;
+        case 'load-bundle':
+          // Load a named bundle of precomputed resources
+          await loadBundle(db, queryCache, path.resolve(BUNDLE_DIR, query.name));
+          res.done();
+          break;
+        default:
+          res.error(`Unrecognized command: ${type}`, 400);
       }
     } catch (err) {
       res.error(err, 500);
-    } finally {
-      res.unlock?.();
     }
 
-    console.log('REQUEST', Math.round(performance.now() - t0));
+    console.log('REQUEST', (performance.now() - t0).toFixed(1));
   };
 }
 
-let locked = false;
-const queue = [];
-
 function httpResponse(res) {
   return {
-    lock() {
-      // if locked, add a promise to the queue
-      // otherwise, grab the lock and proceed
-      return locked
-        ? new Promise(resolve => queue.push(resolve))
-        : (locked = true);
-    },
-    unlock() {
-      locked = queue.length > 0;
-      if (locked) {
-        // resolve the next promise in the queue
-        queue.shift()();
-      }
-    },
-    async stream(iter) {
-      for await (const chunk of iter) {
-        res.write(chunk);
-      }
-      res.end();
+    arrow(data) {
+      res.setHeader('Content-Type', 'application/vnd.apache.arrow.stream');
+      res.end(data);
     },
     json(data) {
+      res.setHeader('Content-Type', 'application/json');
       res.end(JSON.stringify(data));
     },
     done() {
@@ -149,16 +168,11 @@ function httpResponse(res) {
 
 function socketResponse(ws) {
   const STRING = { binary: false, fin: true };
-  const FRAGMENT = { binary: true, fin: false };
-  const DONE = { binary: true, fin: true };
-  const NULL = new Uint8Array(0);
+  const BINARY = { binary: true, fin: true };
 
   return {
-    async stream(iter) {
-      for await (const chunk of iter) {
-        ws.send(chunk, FRAGMENT);
-      }
-      ws.send(NULL, DONE);
+    arrow(data) {
+      ws.send(data, BINARY);
     },
     json(data) {
       ws.send(JSON.stringify(data), STRING);

package/src/index.js

@@ -1,2 +1,6 @@
 export { DuckDB } from './DuckDB.js';
 export { dataServer } from './data-server.js';
+export { loadArrow } from './load/arrow.js';
+export { loadCSV } from './load/csv.js';
+export { loadJSON } from './load/json.js';
+export { loadParquet } from './load/parquet.js';

package/src/load/arrow.js

@@ -0,0 +1,14 @@
+import { readFile } from 'node:fs/promises';
+
+export async function loadArrow(db, tableName, buffer) {
+  const arrowData = ArrayBuffer.isView(buffer) ? buffer : await readFile(buffer);
+  return new Promise((resolve, reject) => {
+    db.con.register_buffer(tableName, [arrowData], true, err => {
+      if (err) {
+        console.error(err);
+        reject(err);
+      }
+      resolve();
+    });
+  });
+}

package/src/load/bundle.js

@@ -0,0 +1,75 @@
+import fs from 'node:fs/promises';
+import path from 'node:path';
+import { cacheKey } from '../Cache.js';
+
+async function retrieve(db, cache, sql, type) {
+  const key = cacheKey(sql, type);
+  const cached = cache.get(key);
+  if (cached) return cached;
+  switch (type) {
+    case 'arrow':
+      return db.arrowBuffer(sql);
+    case 'json':
+      return JSON.stringify(await db.query(sql));
+    default:
+      throw new Error(`Unsupported query type: ${type}`);
+  }
+}
+
+export async function createBundle(db, cache, queries, dir) {
+  const describe_re = /^DESCRIBE /;
+  const pragma_re = /^PRAGMA /;
+  const view_re = /^CREATE( TEMP| TEMPORARY)? VIEW/;
+  const table_re = /^CREATE( TEMP| TEMPORARY)? TABLE( IF NOT EXISTS)? ([^\s]+)/;
+
+  const manifest = { tables: [], queries: [] };
+
+  await fs.mkdir(dir, { recursive: true });
+
+  const querySet = new Set(queries);
+  for (const query of querySet) {
+    const sql = typeof query === 'string' ? query : query.sql;
+    if (query.alias) {
+      const table = query.alias;
+      const file = path.resolve(dir, `${table}.parquet`);
+      await db.exec(`COPY (${sql}) TO '${file}' (FORMAT PARQUET)`);
+      manifest.tables.push(table);
+    } else if (sql.startsWith('CREATE ')) {
+      // table or view
+      if (view_re.test(sql)) continue; // ignore views
+      const table = sql.match(table_re)?.[3];
+      const file = path.resolve(dir, `${table}.parquet`);
+      await db.exec(`${sql}`);
+      await db.exec(`COPY ${table} TO '${file}' (FORMAT PARQUET)`);
+      manifest.tables.push(table);
+    } else if (!pragma_re.test(sql)) {
+      // select query
+      const type = describe_re.test(sql) ? 'json' : 'arrow';
+      const key = cacheKey(sql, type);
+      const result = await retrieve(db, cache, sql, type);
+      await fs.writeFile(path.resolve(dir, key), result);
+      manifest.queries.push(key);
+    }
+  }
+
+  await fs.writeFile(path.resolve(dir, 'bundle.json'), JSON.stringify(manifest, 0, 2));
+  return manifest;
+}
+
+export async function loadBundle(db, cache, dir) {
+  const manifest = JSON.parse(await fs.readFile(path.resolve(dir, 'bundle.json')));
+
+  // load precomputed query results into the cache
+  for (const key of manifest.queries) {
+    const file = path.resolve(dir, key);
+    const json = path.extname(file) === '.json';
+    const data = await fs.readFile(file);
+    cache.set(key, json ? JSON.parse(data) : data);
+  }
+
+  // load precomputed temp tables into the database
+  for (const table of manifest.tables) {
+    const file = path.resolve(dir, `${table}.parquet`);
+    await db.exec(`CREATE TEMP TABLE IF NOT EXISTS ${table} AS SELECT * FROM '${file}'`);
+  }
+}

package/src/load/create-table.js

@@ -0,0 +1,6 @@
+export function createTable(db, name, as, options = {}) {
+  const { temp, replace } = options;
+  const create = `CREATE${replace ? ' OR REPLACE' : ''}`;
+  const type = `${temp ? 'TEMP ' : ''}TABLE${replace ? '' : ' IF NOT EXISTS'}`;
+  return db.exec(`${create} ${type} ${name} AS ${as}`);
+}

package/src/load/csv.js

@@ -0,0 +1,9 @@
+import { createTable } from './create-table.js';
+import { parameters } from './parameters.js';
+
+export function loadCSV(db, tableName, fileName, options = {}) {
+  const { select = ['*'], temp, replace, ...csvOptions } = options;
+  const params = parameters({ auto_detect: true, sample_size: -1, ...csvOptions });
+  const query = `SELECT ${select.join(', ')} FROM read_csv('${fileName}', ${params})`;
+  return createTable(db, tableName, query, { temp, replace });
+}

package/src/load/json.js

@@ -0,0 +1,9 @@
+import { createTable } from './create-table.js';
+import { parameters } from './parameters.js';
+
+export function loadJSON(db, tableName, fileName, options = {}) {
+  const { select = ['*'], temp, replace, ...jsonOptions } = options;
+  const params = parameters({ auto_detect: true, json_format: 'auto', ...jsonOptions });
+  const query = `SELECT ${select.join(', ')} FROM read_json('${fileName}', ${params})`;
+  return createTable(db, tableName, query, { temp, replace });
+}

package/src/load/parameters.js

@@ -0,0 +1,11 @@
+export function parameters(options) {
+  return Object.entries(options)
+    .map(([key, value]) => {
+      const t = typeof value;
+      const v = t === 'boolean' ? String(value)
+        : t === 'string' ? `'${value}'`
+        : value;
+      return `${key}=${v}`;
+    })
+    .join(', ');
+}

package/src/load/parquet.js

@@ -0,0 +1,7 @@
+import { createTable } from './create-table.js';
+
+export function loadParquet(db, tableName, fileName, options = {}) {
+  const { select = ['*'], ...tableOptions } = options;
+  const query = `SELECT ${select.join(', ')} FROM read_parquet('${fileName}')`;
+  return createTable(db, tableName, query, tableOptions);
+}

package/src/merge-buffers.js

@@ -1,10 +1,12 @@
 export function mergeBuffers(buffers) {
-  const len = buffers.reduce((a, b) => a + b.length, 0);
+  const len = buffers.reduce((a, b) => a + (b ? b.length : 0), 0);
   const buf = new Uint8Array(len);
 
   for (let i = 0, offset = 0; i < buffers.length; ++i) {
-    buf.set(buffers[i], offset);
-    offset += buffers[i].length;
+    if (buffers[i]) {
+      buf.set(buffers[i], offset);
+      offset += buffers[i].length;
+    }
   }
 
   return buf;