Merge pull request #111 from duckdb/jray/result-visitors

refactor result extraction to use visitors
duckdb · Jan 18, 2025 · 784a939 · 784a939
2 parents ac600f2 + a3bc3ca
commit 784a939
Show file tree

Hide file tree

Showing 5 changed files with 83 additions and 90 deletions.
diff --git a/api/src/DuckDBDataChunk.ts b/api/src/DuckDBDataChunk.ts
@@ -22,6 +22,12 @@ export class DuckDBDataChunk {
   public get columnCount(): number {
     return duckdb.data_chunk_get_column_count(this.chunk);
   }
+  public get rowCount(): number {
+    return duckdb.data_chunk_get_size(this.chunk);
+  }
+  public set rowCount(count: number) {
+    duckdb.data_chunk_set_size(this.chunk, count);
+  }
   public getColumnVector(columnIndex: number): DuckDBVector {
     if (this.vectors[columnIndex]) {
       return this.vectors[columnIndex];
@@ -33,8 +39,16 @@ export class DuckDBDataChunk {
     this.vectors[columnIndex] = vector;
     return vector;
   }
+  public visitColumnValues(columnIndex: number, visitValue: (value: DuckDBValue, rowIndex?: number, columnIndex?: number) => void) {
+    const vector = this.getColumnVector(columnIndex);
+    for (let rowIndex = 0; rowIndex < vector.itemCount; rowIndex++) {
+      visitValue(vector.getItem(rowIndex), rowIndex, columnIndex);
+    }
+  }
   public getColumnValues(columnIndex: number): DuckDBValue[] {
-    return this.getColumnVector(columnIndex).toArray();
+    const values: DuckDBValue[] = [];
+    this.visitColumnValues(columnIndex, value => values.push(value));
+    return values;
   }
   public setColumnValues(columnIndex: number, values: readonly DuckDBValue[]) {
     const vector = this.getColumnVector(columnIndex);
@@ -46,12 +60,15 @@ export class DuckDBDataChunk {
     }
     vector.flush();
   }
-  public getColumns(): DuckDBValue[][] {
-    const columns: DuckDBValue[][] = [];
+  public visitColumns(visitColumn: (column: DuckDBValue[], columnIndex?: number) => void) {
     const columnCount = this.columnCount;
     for (let columnIndex = 0; columnIndex < columnCount; columnIndex++) {
-      columns.push(this.getColumnValues(columnIndex));
+      visitColumn(this.getColumnValues(columnIndex), columnIndex);
     }
+  }
+  public getColumns(): DuckDBValue[][] {
+    const columns: DuckDBValue[][] = [];
+    this.visitColumns(column => columns.push(column));
     return columns;
   }
   public setColumns(columns: readonly (readonly DuckDBValue[])[]) {
@@ -62,21 +79,32 @@ export class DuckDBDataChunk {
       this.setColumnValues(columnIndex, columns[columnIndex]);
     }
   }
-  public getRows(): DuckDBValue[][] {
-    const rows: DuckDBValue[][] = [];
-    const vectors: DuckDBVector[] = [];
+  public visitColumnMajor(visitValue: (value: DuckDBValue, rowIndex?: number, columnIndex?: number) => void) {
+    const columnCount = this.columnCount;
+    for (let columnIndex = 0; columnIndex < columnCount; columnIndex++) {
+      this.visitColumnValues(columnIndex, visitValue);
+    }
+  }
+  public visitRowValues(rowIndex: number, visitValue: (value: DuckDBValue, rowIndex?: number, columnIndex?: number) => void) {
     const columnCount = this.columnCount;
     for (let columnIndex = 0; columnIndex < columnCount; columnIndex++) {
-      vectors.push(this.getColumnVector(columnIndex));
+      visitValue(this.getColumnVector(columnIndex).getItem(rowIndex), rowIndex, columnIndex);
     }
+  }
+  public getRowValues(rowIndex: number): DuckDBValue[] {
+    const values: DuckDBValue[] = [];
+    this.visitRowValues(rowIndex, value => values.push(value));
+    return values;
+  }
+  public visitRows(visitRow: (row: DuckDBValue[], rowIndex?: number) => void) {
     const rowCount = this.rowCount;
     for (let rowIndex = 0; rowIndex < rowCount; rowIndex++) {
-      const row: DuckDBValue[] = [];
-      for (let columnIndex = 0; columnIndex < columnCount; columnIndex++) {
-        row.push(vectors[columnIndex].getItem(rowIndex));
-      }
-      rows.push(row);
+      visitRow(this.getRowValues(rowIndex), rowIndex);
     }
+  }
+  public getRows(): DuckDBValue[][] {
+    const rows: DuckDBValue[][] = [];
+    this.visitRows(row => rows.push(row));
     return rows;
   }
   public setRows(rows: readonly (readonly DuckDBValue[])[]) {
@@ -90,10 +118,13 @@ export class DuckDBDataChunk {
       vector.flush();
     }
   }
-  public get rowCount(): number {
-    return duckdb.data_chunk_get_size(this.chunk);
-  }
-  public set rowCount(count: number) {
-    duckdb.data_chunk_set_size(this.chunk, count);
+  public visitRowMajor(visitValue: (value: DuckDBValue, rowIndex?: number, columnIndex?: number) => void) {
+    const rowCount = this.rowCount;
+    const columnCount = this.columnCount;
+    for (let rowIndex = 0; rowIndex < rowCount; rowIndex++) {
+      for (let columnIndex = 0; columnIndex < columnCount; columnIndex++) {
+        visitValue(this.getColumnVector(columnIndex).getItem(rowIndex), rowIndex, columnIndex);
+      }
+    }
   }
 }
diff --git a/api/src/DuckDBResult.ts b/api/src/DuckDBResult.ts
@@ -3,8 +3,9 @@ import { DuckDBDataChunk } from './DuckDBDataChunk';
 import { DuckDBLogicalType } from './DuckDBLogicalType';
 import { DuckDBType } from './DuckDBType';
 import { DuckDBTypeId } from './DuckDBTypeId';
-import { DuckDBVector } from './DuckDBVector';
 import { ResultReturnType, StatementType } from './enums';
+import { getColumnsFromChunks } from './getColumnsFromChunks';
+import { getRowsFromChunks } from './getRowsFromChunks';
 import { DuckDBValue } from './values';
 
 export class DuckDBResult {
@@ -78,43 +79,10 @@ export class DuckDBResult {
   }
   public async getColumns(): Promise<DuckDBValue[][]> {
     const chunks = await this.fetchAllChunks();
-    if (chunks.length === 0) {
-      return [];
-    }
-    const firstChunk = chunks[0];
-    const columns: DuckDBValue[][] = [];
-    const columnCount = this.columnCount;
-    for (let columnIndex = 0; columnIndex < columnCount; columnIndex++) {
-      columns.push(firstChunk.getColumnValues(columnIndex));
-    } 
-    for (let chunkIndex = 1; chunkIndex < chunks.length; chunkIndex++) {
-      for (let columnIndex = 0; columnIndex < columnCount; columnIndex++) {
-        const vector = chunks[chunkIndex].getColumnVector(columnIndex);
-        for (let itemIndex = 0; itemIndex < vector.itemCount; itemIndex++) {
-          columns[columnIndex].push(vector.getItem(itemIndex));
-        }
-      }
-    }
-    return columns;
+    return getColumnsFromChunks(chunks);
   }
   public async getRows(): Promise<DuckDBValue[][]> {
     const chunks = await this.fetchAllChunks();
-    const rows: DuckDBValue[][] = [];
-    for (const chunk of chunks) {
-      const chunkVectors: DuckDBVector[] = [];
-      const columnCount = chunk.columnCount;
-      for (let columnIndex = 0; columnIndex < columnCount; columnIndex++) {
-        chunkVectors.push(chunk.getColumnVector(columnIndex));
-      }
-      const rowCount = chunk.rowCount;
-      for (let rowIndex = 0; rowIndex < rowCount; rowIndex++) {
-        const row: DuckDBValue[] = [];
-        for (let columnIndex = 0; columnIndex < columnCount; columnIndex++) {
-          row.push(chunkVectors[columnIndex].getItem(rowIndex));
-        }
-        rows.push(row);
-      }
-    }
-    return rows;
+    return getRowsFromChunks(chunks);
   }
 }
diff --git a/api/src/DuckDBResultReader.ts b/api/src/DuckDBResultReader.ts
@@ -3,8 +3,9 @@ import { DuckDBLogicalType } from './DuckDBLogicalType';
 import { DuckDBResult } from './DuckDBResult';
 import { DuckDBType } from './DuckDBType';
 import { DuckDBTypeId } from './DuckDBTypeId';
-import { DuckDBVector } from './DuckDBVector';
 import { ResultReturnType, StatementType } from './enums';
+import { getColumnsFromChunks } from './getColumnsFromChunks';
+import { getRowsFromChunks } from './getRowsFromChunks';
 import { DuckDBValue } from './values';
 
 interface ChunkSizeRun {
@@ -153,44 +154,11 @@ export class DuckDBResultReader {
   }
 
   public getColumns(): DuckDBValue[][] {
-    if (this.chunks.length === 0) {
-      return [];
-    }
-    const firstChunk = this.chunks[0];
-    const columns: DuckDBValue[][] = [];
-    const columnCount = this.columnCount;
-    for (let columnIndex = 0; columnIndex < columnCount; columnIndex++) {
-      columns.push(firstChunk.getColumnValues(columnIndex));
-    } 
-    for (let chunkIndex = 1; chunkIndex < this.chunks.length; chunkIndex++) {
-      for (let columnIndex = 0; columnIndex < columnCount; columnIndex++) {
-        const vector = this.chunks[chunkIndex].getColumnVector(columnIndex);
-        for (let itemIndex = 0; itemIndex < vector.itemCount; itemIndex++) {
-          columns[columnIndex].push(vector.getItem(itemIndex));
-        }
-      }
-    }
-    return columns;
+    return getColumnsFromChunks(this.chunks);
   }
 
   public getRows(): DuckDBValue[][] {
-    const rows: DuckDBValue[][] = [];
-    for (const chunk of this.chunks) {
-      const chunkVectors: DuckDBVector[] = [];
-      const columnCount = chunk.columnCount;
-      for (let columnIndex = 0; columnIndex < columnCount; columnIndex++) {
-        chunkVectors.push(chunk.getColumnVector(columnIndex));
-      }
-      const rowCount = chunk.rowCount;
-      for (let rowIndex = 0; rowIndex < rowCount; rowIndex++) {
-        const row: DuckDBValue[] = [];
-        for (let columnIndex = 0; columnIndex < columnCount; columnIndex++) {
-          row.push(chunkVectors[columnIndex].getItem(rowIndex));
-        }
-        rows.push(row);
-      }
-    }
-    return rows;
+    return getRowsFromChunks(this.chunks);
   }
 
 }
diff --git a/api/src/getColumnsFromChunks.ts b/api/src/getColumnsFromChunks.ts
@@ -0,0 +1,16 @@
+import { DuckDBDataChunk } from './DuckDBDataChunk';
+import { DuckDBValue } from './values';
+
+export function getColumnsFromChunks(chunks: readonly DuckDBDataChunk[]): DuckDBValue[][] {
+  const columns: DuckDBValue[][] = [];
+  if (chunks.length === 0) {
+    return columns;
+  }
+  chunks[0].visitColumns(column => columns.push(column));
+  for (let chunkIndex = 1; chunkIndex < chunks.length; chunkIndex++) {
+    for (let columnIndex = 0; columnIndex < columns.length; columnIndex++) {
+      chunks[chunkIndex].visitColumnValues(columnIndex, value => columns[columnIndex].push(value));
+    }
+  }
+  return columns;
+}
diff --git a/api/src/getRowsFromChunks.ts b/api/src/getRowsFromChunks.ts
@@ -0,0 +1,10 @@
+import { DuckDBDataChunk } from './DuckDBDataChunk';
+import { DuckDBValue } from './values';
+
+export function getRowsFromChunks(chunks: readonly DuckDBDataChunk[]): DuckDBValue[][] {
+  const rows: DuckDBValue[][] = [];
+  for (const chunk of chunks) {
+    chunk.visitRows(row => rows.push(row));
+  }
+  return rows;
+}