MotherDuck unwrapping (#1658)

* Handle MotherDuck types better * Fix decimal
malloydata · Mar 14, 2024 · 30e246b · 30e246b
1 parent 47007fc
commit 30e246b
Show file tree

Hide file tree

Showing 2 changed files with 66 additions and 12 deletions.
diff --git a/packages/malloy-db-duckdb/src/duckdb_wasm_connection.ts b/packages/malloy-db-duckdb/src/duckdb_wasm_connection.ts
@@ -69,12 +69,6 @@ export const unwrapArrow = (value: unknown): any => {
       return Number(obj[Symbol.toPrimitive]());
     } else if (Array.isArray(value)) {
       return value.map(unwrapArrow);
-    } else if (obj['microseconds'] && obj['timezone'] === null) {
-      // Convert epoch µs to ms
-      return Number(obj['microseconds']) / 1000;
-    } else if (obj['days']) {
-      // Convert epoch day to Date
-      return new Date(obj['days'] * 8.64e7);
     } else {
       // eslint-disable-next-line @typescript-eslint/no-explicit-any
       const result: Record<string | symbol, any> = {};

diff --git a/packages/malloy-db-duckdb/src/duckdb_wasm_connection_browser.ts b/packages/malloy-db-duckdb/src/duckdb_wasm_connection_browser.ts
@@ -22,12 +22,70 @@
  */
 
 import * as duckdb from '@duckdb/duckdb-wasm';
+import {DuckDBWASMConnection as DuckDBWASMConnectionBase} from './duckdb_wasm_connection';
 import {
-  DuckDBWASMConnection as DuckDBWASMConnectionBase,
-  unwrapArrow,
-} from './duckdb_wasm_connection';
-import {MDConnection} from '@motherduck/wasm-client';
-import {QueryDataRow} from '@malloydata/malloy';
+  // DuckDBBlob,
+  DuckDBDate,
+  DuckDBDecimal,
+  // DuckDBInterval,
+  DuckDBList,
+  // DuckDBMap,
+  DuckDBRow,
+  DuckDBStruct,
+  DuckDBTime,
+  DuckDBTimestampMicroseconds,
+  DuckDBTimestampMilliseconds,
+  DuckDBTimestampNanoseconds,
+  DuckDBTimestampSeconds,
+  DuckDBValue,
+  MDConnection,
+} from '@motherduck/wasm-client';
+import {QueryDataRow, QueryValue} from '@malloydata/malloy';
+
+function unwrapMotherDuck(value: DuckDBValue) {
+  let result: QueryValue = null;
+  if (value !== null && typeof value === 'object') {
+    if (value instanceof DuckDBDate) {
+      result = new Date(value.days * 8.64e7);
+    } else if (value instanceof DuckDBDecimal) {
+      result = Number(value.scaledValue) / Math.pow(10, value.scale);
+    } else if (value instanceof DuckDBTime) {
+      result = new Date(Number(value.microseconds) / 1000);
+    } else if (value instanceof DuckDBTimestampMicroseconds) {
+      result = new Date(Number(value.microseconds) / 1000);
+    } else if (value instanceof DuckDBTimestampMilliseconds) {
+      result = new Date(Number(value.milliseconds));
+    } else if (value instanceof DuckDBTimestampNanoseconds) {
+      result = new Date(Number(value.nanoseconds) / 1000 / 1000);
+    } else if (value instanceof DuckDBTimestampSeconds) {
+      result = new Date(Number(value.seconds) * 1000);
+    } else if (value instanceof DuckDBStruct) {
+      const struct: QueryValue = {};
+      for (const structEntry of value.entries) {
+        struct[structEntry.key] = unwrapMotherDuck(structEntry.value);
+      }
+      result = struct;
+    } else if (value instanceof DuckDBList) {
+      result = value.values.map(unwrapMotherDuck);
+    } else {
+      result = value.toString();
+    }
+  } else if (typeof value === 'bigint') {
+    result = Number(value);
+  } else {
+    result = value;
+  }
+  return result;
+}
+
+function unwrapMotherDuckRow(row: DuckDBRow) {
+  const result: QueryDataRow = {};
+  for (const entry of Object.entries(row)) {
+    const [key, value] = entry;
+    result[key] = unwrapMotherDuck(value);
+  }
+  return result;
+}
 
 export class DuckDBWASMConnection extends DuckDBWASMConnectionBase {
   protected _mdConnection: MDConnection | null = null;
@@ -93,7 +151,9 @@ export class DuckDBWASMConnection extends DuckDBWASMConnectionBase {
         if (queryId) {
           const result = await connection.evaluateQueuedQuery(queryId);
           if (result?.data) {
-            const rows = unwrapArrow(result.data.toRows());
+            const rows = result.data
+              .toRows()
+              .map(row => unwrapMotherDuckRow(row));
             const totalRows = result.data.rowCount;
             return {
               rows,