Apache Arrow for DuckDBClient

Author: mbostock

src/arrow.mjs

@@ -1,3 +1,17 @@
+// Returns true if the vaue is an Apache Arrow table. This uses a “duck” test
+// (instead of strict instanceof) because we want it to work with a range of
+// Apache Arrow versions at least 7.0.0 or above.
+// https://arrow.apache.org/docs/7.0/js/classes/Arrow_dom.Table.html
+export function isArrowTable(value) {
+  return (
+    value &&
+    typeof value.getChild === "function" &&
+    typeof value.toArray === "function" &&
+    value.schema &&
+    Array.isArray(value.schema.fields)
+  );
+}
+
 export function getArrowTableSchema(table) {
   return table.schema.fields.map(getArrowFieldSchema);
 }

src/duckdb.mjs

@@ -1,6 +1,7 @@
-import {getArrowTableSchema} from "./arrow.mjs";
+import {getArrowTableSchema, isArrowTable} from "./arrow.mjs";
 import {arrow9 as arrow, duckdb} from "./dependencies.mjs";
 import {FileAttachment} from "./fileAttachment.mjs";
+import {cdn} from "./require.mjs";
 
 // Adapted from https://observablehq.com/@cmudig/duckdb-client
 // Copyright 2021 CMU Data Interaction Group
@@ -31,9 +32,6 @@ import {FileAttachment} from "./fileAttachment.mjs";
 // ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 // POSSIBILITY OF SUCH DAMAGE.
 
-// TODO Allow this to be overridden using the Library’s resolver.
-const cdn = "https://cdn.observableusercontent.com/npm/";
-
 export class DuckDBClient {
   constructor(db) {
     Object.defineProperties(this, {
@@ -125,16 +123,22 @@ export class DuckDBClient {
     await db.open(config);
     await Promise.all(
       Object.entries(sources).map(async ([name, source]) => {
-        if ("array" in source) { // array + options
-          const {array, ...options} = source;
-          await insertArray(db, name, array, options);
+        if (source instanceof FileAttachment) { // bare file
+          await insertFile(db, name, source);
+        } else if (isArrowTable(source)) { // bare arrow table
+          await insertArrowTable(db, name, source);
+        } else if (Array.isArray(source)) { // bare array of objects
+          await insertArray(db, name, source);
+        } else if ("data" in source) { // data + options
+          const {data, ...options} = source;
+          if (isArrowTable(data)) {
+            await insertArrowTable(db, name, data, options);
+          } else {
+            await insertArray(db, name, data, options);
+          }
         } else if ("file" in source) { // file + options
           const {file, ...options} = source;
           await insertFile(db, name, file, options);
-        } else if (source instanceof FileAttachment) { // bare file
-          await insertFile(db, name, source);
-        } else if (Array.isArray(source)) { // bare data
-          await insertArray(db, name, source);
         } else {
           throw new Error(`invalid source: ${source}`);
         }
@@ -156,36 +160,40 @@ async function insertFile(database, name, file, options) {
   try {
     switch (file.mimeType) {
       case "text/csv":
-        await connection.insertCSVFromPath(file.name, {
+        return await connection.insertCSVFromPath(file.name, {
           name,
           schema: "main",
           ...options
         });
-        break;
       case "application/json":
-        await connection.insertJSONFromPath(file.name, {
+        return await connection.insertJSONFromPath(file.name, {
           name,
           schema: "main",
           ...options
         });
-        break;
       default:
-        if (file.name.endsWith(".parquet")) {
-          await connection.query(
+        if (/\.arrow$/i.test(file.name)) {
+          const buffer = new Uint8Array(await file.arrayBuffer());
+          return await connection.insertArrowFromIPCStream(buffer, {
+            name,
+            schema: "main",
+            ...options
+          });
+        }
+        if (/\.parquet$/i.test(file.name)) {
+          return await connection.query(
             `CREATE VIEW '${name}' AS SELECT * FROM parquet_scan('${file.name}')`
           );
-        } else {
-          throw new Error(`unknown file type: ${file.mimeType}`);
         }
+        throw new Error(`unknown file type: ${file.mimeType}`);
     }
   } finally {
     await connection.close();
   }
 }
 
-async function insertArray(database, name, array, options) {
+async function insertArrowTable(database, name, table, options) {
   const arrow = await loadArrow();
-  const table = arrow.tableFromJSON(array);
   const buffer = arrow.tableToIPC(table);
   const connection = await database.connect();
   try {
@@ -199,6 +207,12 @@ async function insertArray(database, name, array, options) {
   }
 }
 
+async function insertArray(database, name, array, options) {
+  const arrow = await loadArrow();
+  const table = arrow.tableFromJSON(array);
+  return await insertArrowTable(database, name, table, options);
+}
+
 async function createDuckDB() {
   const duck = await import(`${cdn}${duckdb.resolve()}`);
   const bundle = await duck.selectBundle({

src/fileAttachment.mjs

@@ -1,6 +1,6 @@
 import {autoType, csvParse, csvParseRows, tsvParse, tsvParseRows} from "d3-dsv";
-import {arrow4, jszip, exceljs} from "./dependencies.mjs";
-import {requireDefault} from "./require.mjs";
+import {arrow4, arrow9, jszip, exceljs} from "./dependencies.mjs";
+import {cdn, requireDefault} from "./require.mjs";
 import {SQLiteDatabaseClient} from "./sqlite.mjs";
 import {Workbook} from "./xlsx.mjs";
 
@@ -56,9 +56,18 @@ export class AbstractFile {
       i.src = url;
     });
   }
-  async arrow() {
-    const [Arrow, response] = await Promise.all([requireDefault(arrow4.resolve()), remote_fetch(this)]); // TODO upgrade to apache-arrow@9
-    return Arrow.Table.from(response);
+  async arrow({version = 4} = {}) {
+    switch (version) {
+      case 4: {
+        const [Arrow, response] = await Promise.all([requireDefault(arrow4.resolve()), remote_fetch(this)]);
+        return Arrow.Table.from(response);
+      }
+      case 9: {
+        const [Arrow, response] = await Promise.all([import(`${cdn}${arrow9.resolve()}`), remote_fetch(this)]);
+        return Arrow.tableFromIPC(response);
+      }
+      default: throw new Error(`unsupported arrow version: ${version}`);
+    }
   }
   async sqlite() {
     return SQLiteDatabaseClient.open(remote_fetch(this));

src/index.mjs

@@ -1,3 +1,3 @@
 export {default as FileAttachments, AbstractFile} from "./fileAttachment.mjs";
 export {default as Library} from "./library.mjs";
-export {makeQueryTemplate, loadDataSource, arrayIsPrimitive, isDataArray, isDatabaseClient} from "./table.mjs";
+export {makeQueryTemplate, loadDataSource, arrayIsPrimitive, isArrowTable, isDataArray, isDatabaseClient} from "./table.mjs";

src/require.mjs

@@ -1,5 +1,8 @@
 import {require as initialRequire, requireFrom} from "d3-require";
 
+// TODO Allow this to be overridden using the Library’s resolver.
+export const cdn = "https://cdn.observableusercontent.com/npm/";
+
 export let requireDefault = initialRequire;
 
 export function setDefaultRequire(require) {