@uwdata/mosaic-duckdb 0.0.1

package/README.md

@@ -0,0 +1,3 @@
+# @mosaic/duckdb
+
+A Node.js + DuckDB data server that supports transfer of [Apache Arrow](https://arrow.apache.org/) and JSON data over either Web Sockets or HTTP.

package/bin/csv2arrow.js

@@ -0,0 +1,29 @@
+#! /usr/bin/env node
+import { DuckDB } from '../src/index.js';
+import { createWriteStream } from 'fs';
+
+const db = new DuckDB();
+
+// load CSV into duckdb
+await db.csv('data', process.argv[2]);
+
+// get output stream of arrow bytes
+const stream = await db.arrowStream('SELECT * FROM data');
+
+ // determine the output stream
+const output = process.argv[3]
+  ? createWriteStream(process.argv[3])
+  : process.stdout;
+
+// set up error handling
+output.on('error', (error) => {
+  console.error(`File write error: ${error.message}`);
+});
+
+// write arrow bytes to output
+for await (const chunk of stream) {
+  output.write(chunk);
+}
+
+// finish
+output.end(new Uint8Array(4));

package/package.json

@@ -0,0 +1,30 @@
+{
+  "name": "@uwdata/mosaic-duckdb",
+  "version": "0.0.1",
+  "description": "A Node-based DuckDB data server.",
+  "keywords": [
+    "duckdb",
+    "server",
+    "node",
+    "arrow",
+    "mosaic"
+  ],
+  "license": "BSD-3-Clause",
+  "author": "Jeffrey Heer (http://idl.cs.washington.edu)",
+  "type": "module",
+  "main": "src/index.js",
+  "module": "src/index.js",
+  "repository": {
+    "type": "git",
+    "url": "https://github.com/uwdata/mosaic.git"
+  },
+  "scripts": {
+    "lint": "eslint src test --ext .js",
+    "test": "mocha 'test/**/*-test.js'",
+    "prepublishOnly": "npm run test && npm run lint"
+  },
+  "dependencies": {
+    "duckdb": "~0.6.1",
+    "ws": "^8.12.0"
+  }
+}

package/src/DuckDB.js

@@ -0,0 +1,158 @@
+import duckdb from 'duckdb';
+import { readFile } from 'node:fs/promises';
+import { mergeBuffers } from './merge-buffers.js';
+
+const CONFIG = [
+  `PRAGMA temp_directory='./duckdb.tmp'`,
+  `INSTALL arrow`,
+  `INSTALL httpfs`,
+  `LOAD arrow`,
+  `LOAD httpfs`
+];
+
+export class DuckDB {
+  constructor(path = ':memory:') {
+    this.db = new duckdb.Database(path);
+    this.con = this.db.connect();
+    this.exec(CONFIG.join(';\n'));
+  }
+
+  close() {
+    return new Promise((resolve, reject) => {
+      this.db.close((err) => {
+        if (err) {
+          reject(err);
+        } else {
+          resolve(this);
+        }
+      });
+    });
+  }
+
+  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));
+  }
+
+  exec(sql) {
+    return new Promise((resolve, reject) => {
+      this.con.exec(sql, (err) => {
+        if (err) {
+          reject(err);
+        } else {
+          resolve(this);
+        }
+      });
+    });
+  }
+
+  query(sql) {
+    return new Promise((resolve, reject) => {
+      this.con.all(sql, (err, result) => {
+        if (err) {
+          reject(err);
+        } else {
+          resolve(result);
+        }
+      });
+    });
+  }
+
+  arrowBuffer(sql) {
+    return new Promise((resolve, reject) => {
+      this.con.arrowIPCAll(sql, (err, result) => {
+        if (err) {
+          reject(err);
+        } else {
+          resolve(mergeBuffers(result));
+        }
+      });
+    });
+  }
+
+  arrowStream(sql) {
+    return this.con.arrowIPCStream(sql);
+  }
+}
+
+export class DuckDBStatement {
+  constructor(statement) {
+    this.statement = statement;
+  }
+
+  finalize() {
+    this.statement.finalize();
+  }
+
+  run(params) {
+    this.statement.run(...params);
+  }
+
+  exec(params) {
+    return new Promise((resolve, reject) => {
+      this.statement.run(...params, (err) => {
+        if (err) {
+          reject(err);
+        } else {
+          resolve(this);
+        }
+      });
+    });
+  }
+
+  query(params) {
+    return new Promise((resolve, reject) => {
+      this.statement.all(...params, (err, result) => {
+        if (err) {
+          reject(err);
+        } else {
+          resolve(result);
+        }
+      });
+    });
+  }
+
+  arrowBuffer(params) {
+    return new Promise((resolve, reject) => {
+      this.con.arrowIPCAll(...params, (err, result) => {
+        if (err) {
+          reject(err);
+        } else {
+          resolve(mergeBuffers(result));
+        }
+      });
+    });
+  }
+
+  arrowStream(params) {
+    return this.statement.arrowIPCStream(...params);
+  }
+}

package/src/data-server.js

@@ -0,0 +1,174 @@
+import http from 'node:http';
+import url from 'node:url';
+import { WebSocketServer } from 'ws';
+
+export function dataServer(db, {
+  rest = true,
+  socket = true,
+  port = 3000
+} = {}) {
+  const handleQuery = queryHandler(db);
+  const app = createHTTPServer(handleQuery, rest);
+  if (socket) createSocketServer(app, handleQuery);
+
+  app.listen(port);
+  console.log(`Data server running on port ${port}`);
+  if (rest) console.log(`  http://localhost:${port}/`);
+  if (socket) console.log(`  ws://localhost:${port}/`);
+}
+
+function createHTTPServer(handleQuery, rest) {
+  return http.createServer((req, resp) => {
+    const res = httpResponse(resp);
+    if (!rest) {
+      res.done();
+      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');
+    resp.setHeader('Access-Control-Allow-Headers', '*');
+    resp.setHeader('Access-Control-Max-Age', 2592000);
+
+    switch (req.method) {
+      case 'OPTIONS':
+        res.done();
+        break;
+      case 'GET':
+        handleQuery(res, url.parse(req.url, true).query);
+        break;
+      case 'POST': {
+        const chunks = [];
+        req.on('error', err => res.error(err, 500));
+        req.on('data', chunk => chunks.push(chunk));
+        req.on('end', () => handleQuery(res, Buffer.concat(chunks)));
+        break;
+      }
+      default:
+        res.error(`Unsupported HTTP method: ${req.method}`, 400);
+    }
+  });
+}
+
+function createSocketServer(server, handleQuery) {
+  const wss = new WebSocketServer({ server });
+
+  wss.on('connection', socket => {
+    const res = socketResponse(socket);
+    socket.on('message', data => handleQuery(res, data));
+  });
+}
+
+function queryHandler(db) {
+  return async (res, data) => {
+    const t0 = performance.now();
+
+    // parse incoming query
+    let query;
+    try {
+      query = JSON.parse(data);
+    } catch (err) {
+      res.error(err, 400);
+      return;
+    }
+
+    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?.();
+
+      // 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:
+          // JSON response format
+          res.json(await db.query(sql));
+      }
+    } catch (err) {
+      res.error(err, 500);
+    } finally {
+      res.unlock?.();
+    }
+
+    console.log('REQUEST', Math.round(performance.now() - t0));
+  };
+}
+
+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();
+    },
+    json(data) {
+      res.end(JSON.stringify(data));
+    },
+    done() {
+      res.writeHead(200);
+      res.end();
+    },
+    error(err, code) {
+      console.error(err);
+      res.writeHead(code);
+      res.end();
+    }
+  }
+}
+
+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);
+
+  return {
+    async stream(iter) {
+      for await (const chunk of iter) {
+        ws.send(chunk, FRAGMENT);
+      }
+      ws.send(NULL, DONE);
+    },
+    json(data) {
+      ws.send(JSON.stringify(data), STRING);
+    },
+    done() {
+      this.json({});
+    },
+    error(err) {
+      console.error(err);
+      this.json({ error: String(err) });
+    }
+  };
+}

package/src/index.js

@@ -0,0 +1,2 @@
+export { DuckDB } from './DuckDB.js';
+export { dataServer } from './data-server.js';

package/src/merge-buffers.js

@@ -0,0 +1,11 @@
+export function mergeBuffers(buffers) {
+  const len = buffers.reduce((a, b) => a + b.length, 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;
+  }
+
+  return buf;
+}