Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

row and column objects #112

Merged
merged 1 commit into from
Jan 18, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
69 changes: 55 additions & 14 deletions api/src/DuckDBDataChunk.ts
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,13 @@ export class DuckDBDataChunk {
constructor(chunk: duckdb.DataChunk) {
this.chunk = chunk;
}
public static create(types: readonly DuckDBType[], rowCount?: number): DuckDBDataChunk {
const chunk = new DuckDBDataChunk(duckdb.create_data_chunk(types.map(t => t.toLogicalType().logical_type)));
public static create(
types: readonly DuckDBType[],
rowCount?: number
): DuckDBDataChunk {
const chunk = new DuckDBDataChunk(
duckdb.create_data_chunk(types.map((t) => t.toLogicalType().logical_type))
);
if (rowCount != undefined) {
chunk.rowCount = rowCount;
}
Expand Down Expand Up @@ -39,15 +44,22 @@ export class DuckDBDataChunk {
this.vectors[columnIndex] = vector;
return vector;
}
public visitColumnValues(columnIndex: number, visitValue: (value: DuckDBValue, rowIndex?: number, columnIndex?: number) => void) {
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[] {
const values: DuckDBValue[] = [];
this.visitColumnValues(columnIndex, value => values.push(value));
this.visitColumnValues(columnIndex, (value) => values.push(value));
return values;
}
public setColumnValues(columnIndex: number, values: readonly DuckDBValue[]) {
Expand All @@ -60,15 +72,17 @@ export class DuckDBDataChunk {
}
vector.flush();
}
public visitColumns(visitColumn: (column: DuckDBValue[], columnIndex?: number) => void) {
public visitColumns(
visitColumn: (column: DuckDBValue[], columnIndex: number) => void
) {
const columnCount = this.columnCount;
for (let columnIndex = 0; columnIndex < columnCount; columnIndex++) {
visitColumn(this.getColumnValues(columnIndex), columnIndex);
}
}
public getColumns(): DuckDBValue[][] {
const columns: DuckDBValue[][] = [];
this.visitColumns(column => columns.push(column));
this.visitColumns((column) => columns.push(column));
return columns;
}
public setColumns(columns: readonly (readonly DuckDBValue[])[]) {
Expand All @@ -79,32 +93,49 @@ export class DuckDBDataChunk {
this.setColumnValues(columnIndex, columns[columnIndex]);
}
}
public visitColumnMajor(visitValue: (value: DuckDBValue, rowIndex?: number, columnIndex?: number) => void) {
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) {
public visitRowValues(
rowIndex: number,
visitValue: (
value: DuckDBValue,
rowIndex: number,
columnIndex: number
) => void
) {
const columnCount = this.columnCount;
for (let columnIndex = 0; columnIndex < columnCount; columnIndex++) {
visitValue(this.getColumnVector(columnIndex).getItem(rowIndex), rowIndex, columnIndex);
visitValue(
this.getColumnVector(columnIndex).getItem(rowIndex),
rowIndex,
columnIndex
);
}
}
public getRowValues(rowIndex: number): DuckDBValue[] {
const values: DuckDBValue[] = [];
this.visitRowValues(rowIndex, value => values.push(value));
this.visitRowValues(rowIndex, (value) => values.push(value));
return values;
}
public visitRows(visitRow: (row: DuckDBValue[], rowIndex?: number) => void) {
public visitRows(visitRow: (row: DuckDBValue[], rowIndex: number) => void) {
const rowCount = this.rowCount;
for (let rowIndex = 0; rowIndex < rowCount; rowIndex++) {
visitRow(this.getRowValues(rowIndex), rowIndex);
}
}
public getRows(): DuckDBValue[][] {
const rows: DuckDBValue[][] = [];
this.visitRows(row => rows.push(row));
this.visitRows((row) => rows.push(row));
return rows;
}
public setRows(rows: readonly (readonly DuckDBValue[])[]) {
Expand All @@ -118,12 +149,22 @@ export class DuckDBDataChunk {
vector.flush();
}
}
public visitRowMajor(visitValue: (value: DuckDBValue, rowIndex?: number, columnIndex?: number) => void) {
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);
visitValue(
this.getColumnVector(columnIndex).getItem(rowIndex),
rowIndex,
columnIndex
);
}
}
}
Expand Down
26 changes: 26 additions & 0 deletions api/src/DuckDBResult.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,8 @@ import { DuckDBType } from './DuckDBType';
import { DuckDBTypeId } from './DuckDBTypeId';
import { ResultReturnType, StatementType } from './enums';
import { getColumnsFromChunks } from './getColumnsFromChunks';
import { getColumnsObjectFromChunks } from './getColumnsObjectFromChunks';
import { getRowObjectsFromChunks } from './getRowObjectsFromChunks';
import { getRowsFromChunks } from './getRowsFromChunks';
import { DuckDBValue } from './values';

Expand Down Expand Up @@ -33,6 +35,22 @@ export class DuckDBResult {
}
return columnNames;
}
public deduplicatedColumnNames(): string[] {
const outputColumnNames: string[] = [];
const columnCount = this.columnCount;
const columnNameCount: { [columnName: string]: number } = {};
for (let columnIndex = 0; columnIndex < columnCount; columnIndex++) {
const inputColumnName = this.columnName(columnIndex);
const nameCount = (columnNameCount[inputColumnName] || 0) + 1;
columnNameCount[inputColumnName] = nameCount;
if (nameCount > 1) {
outputColumnNames.push(`${inputColumnName}:${nameCount - 1}`);
} else {
outputColumnNames.push(inputColumnName);
}
}
return outputColumnNames;
}
public columnTypeId(columnIndex: number): DuckDBTypeId {
return duckdb.column_type(
this.result,
Expand Down Expand Up @@ -81,8 +99,16 @@ export class DuckDBResult {
const chunks = await this.fetchAllChunks();
return getColumnsFromChunks(chunks);
}
public async getColumnsObject(): Promise<Record<string, DuckDBValue[]>> {
const chunks = await this.fetchAllChunks();
return getColumnsObjectFromChunks(chunks, this.deduplicatedColumnNames());
}
public async getRows(): Promise<DuckDBValue[][]> {
const chunks = await this.fetchAllChunks();
return getRowsFromChunks(chunks);
}
public async getRowObjects(): Promise<Record<string, DuckDBValue>[]> {
const chunks = await this.fetchAllChunks();
return getRowObjectsFromChunks(chunks, this.deduplicatedColumnNames());
}
}
19 changes: 16 additions & 3 deletions api/src/DuckDBResultReader.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,8 @@ import { DuckDBType } from './DuckDBType';
import { DuckDBTypeId } from './DuckDBTypeId';
import { ResultReturnType, StatementType } from './enums';
import { getColumnsFromChunks } from './getColumnsFromChunks';
import { getColumnsObjectFromChunks } from './getColumnsObjectFromChunks';
import { getRowObjectsFromChunks } from './getRowObjectsFromChunks';
import { getRowsFromChunks } from './getRowsFromChunks';
import { DuckDBValue } from './values';

Expand Down Expand Up @@ -44,6 +46,9 @@ export class DuckDBResultReader {
public columnNames(): string[] {
return this.result.columnNames();
}
public deduplicatedColumnNames(): string[] {
return this.result.deduplicatedColumnNames();
}
public columnTypeId(columnIndex: number): DuckDBTypeId {
return this.result.columnTypeId(columnIndex);
}
Expand Down Expand Up @@ -99,10 +104,10 @@ export class DuckDBResultReader {
}
// We didn't find our row. It must have been out of range.
throw Error(
`Row index ${rowIndex} requested, but only ${this.currentRowCount_} row have been read so far.`,
`Row index ${rowIndex} requested, but only ${this.currentRowCount_} row have been read so far.`
);
}

/** Read all rows. */
public async readAll(): Promise<void> {
return this.fetchChunks();
Expand All @@ -121,7 +126,8 @@ export class DuckDBResultReader {
while (
!(
this.done_ ||
(targetRowCount !== undefined && this.currentRowCount_ >= targetRowCount)
(targetRowCount !== undefined &&
this.currentRowCount_ >= targetRowCount)
)
) {
const chunk = await this.result.fetchChunk();
Expand Down Expand Up @@ -157,8 +163,15 @@ export class DuckDBResultReader {
return getColumnsFromChunks(this.chunks);
}

public getColumnsObject(): Record<string, DuckDBValue[]> {
return getColumnsObjectFromChunks(this.chunks, this.deduplicatedColumnNames());
}

public getRows(): DuckDBValue[][] {
return getRowsFromChunks(this.chunks);
}

public getRowObjecs(): Record<string, DuckDBValue>[] {
return getRowObjectsFromChunks(this.chunks, this.deduplicatedColumnNames());
}
}
10 changes: 7 additions & 3 deletions api/src/getColumnsFromChunks.ts
Original file line number Diff line number Diff line change
@@ -1,15 +1,19 @@
import { DuckDBDataChunk } from './DuckDBDataChunk';
import { DuckDBValue } from './values';

export function getColumnsFromChunks(chunks: readonly DuckDBDataChunk[]): DuckDBValue[][] {
export function getColumnsFromChunks(
chunks: readonly DuckDBDataChunk[]
): DuckDBValue[][] {
const columns: DuckDBValue[][] = [];
if (chunks.length === 0) {
return columns;
}
chunks[0].visitColumns(column => columns.push(column));
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));
chunks[chunkIndex].visitColumnValues(columnIndex, (value) =>
columns[columnIndex].push(value)
);
}
}
return columns;
Expand Down
24 changes: 24 additions & 0 deletions api/src/getColumnsObjectFromChunks.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
import { DuckDBDataChunk } from './DuckDBDataChunk';
import { DuckDBValue } from './values';

export function getColumnsObjectFromChunks(
chunks: readonly DuckDBDataChunk[],
columnNames: readonly string[],
): Record<string, DuckDBValue[]> {
const columnsObject: Record<string, DuckDBValue[]> = {};
for (const columnName of columnNames) {
columnsObject[columnName] = [];
}
if (chunks.length === 0) {
return columnsObject;
}
const columnCount = chunks[0].columnCount;
for (let chunkIndex = 0; chunkIndex < chunks.length; chunkIndex++) {
for (let columnIndex = 0; columnIndex < columnCount; columnIndex++) {
chunks[chunkIndex].visitColumnValues(columnIndex, (value) =>
columnsObject[columnNames[columnIndex]].push(value)
);
}
}
return columnsObject;
}
20 changes: 20 additions & 0 deletions api/src/getRowObjectsFromChunks.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
import { DuckDBDataChunk } from './DuckDBDataChunk';
import { DuckDBValue } from './values';

export function getRowObjectsFromChunks(
chunks: readonly DuckDBDataChunk[],
columnNames: readonly string[]
): Record<string, DuckDBValue>[] {
const rowObjects: Record<string, DuckDBValue>[] = [];
for (const chunk of chunks) {
const rowCount = chunk.rowCount;
for (let rowIndex = 0; rowIndex < rowCount; rowIndex++) {
const rowObject: Record<string, DuckDBValue> = {};
chunk.visitRowValues(rowIndex, (value, _, columnIndex) => {
rowObject[columnNames[columnIndex]] = value;
});
rowObjects.push(rowObject);
}
}
return rowObjects;
}
6 changes: 4 additions & 2 deletions api/src/getRowsFromChunks.ts
Original file line number Diff line number Diff line change
@@ -1,10 +1,12 @@
import { DuckDBDataChunk } from './DuckDBDataChunk';
import { DuckDBValue } from './values';

export function getRowsFromChunks(chunks: readonly DuckDBDataChunk[]): DuckDBValue[][] {
export function getRowsFromChunks(
chunks: readonly DuckDBDataChunk[]
): DuckDBValue[][] {
const rows: DuckDBValue[][] = [];
for (const chunk of chunks) {
chunk.visitRows(row => rows.push(row));
chunk.visitRows((row) => rows.push(row));
}
return rows;
}
20 changes: 20 additions & 0 deletions api/test/api.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1028,6 +1028,26 @@ describe('api', () => {
]);
});
});
test('row and column objects', async () => {
await withConnection(async (connection) => {
const reader = await connection.runAndReadAll(
'select i::int as a, i::int + 10 as b, (i + 100)::varchar as a from range(3) t(i)'
);
assert.deepEqual(reader.columnNames(), ['a', 'b', 'a']);
assert.deepEqual(reader.deduplicatedColumnNames(), ['a', 'b', 'a:1']);
assert.deepEqual(reader.columnTypes(), [INTEGER, INTEGER, VARCHAR]);
assert.deepEqual(reader.getRowObjecs(), [
{ 'a': 0, 'b': 10, 'a:1': '100' },
{ 'a': 1, 'b': 11, 'a:1': '101' },
{ 'a': 2, 'b': 12, 'a:1': '102' },
]);
assert.deepEqual(reader.getColumnsObject(), {
'a': [0, 1, 2],
'b': [10, 11, 12],
'a:1': ['100', '101', '102'],
});
});
});
test('result reader', async () => {
await withConnection(async (connection) => {
const reader = await connection.runAndReadAll(
Expand Down
Loading