Skip to content

Commit 568d95c

Browse files
jraymakersjraymakers
committed
api: implement varint, fix timestamp_tz vector type
1 parent 3e8588b commit 568d95c

File tree

2 files changed

+101
-29
lines changed

2 files changed

+101
-29
lines changed

‎api/src/DuckDBVector.ts

Lines changed: 93 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
1-
import os from 'os';
21
import duckdb from '@duckdb/node-bindings';
2+
import os from 'os';
33
import { DuckDBLogicalType } from './DuckDBLogicalType';
44
import {
55
DuckDBArrayType,
@@ -25,6 +25,7 @@ import {
2525
DuckDBTimestampMillisecondsType,
2626
DuckDBTimestampNanosecondsType,
2727
DuckDBTimestampSecondsType,
28+
DuckDBTimestampTZType,
2829
DuckDBTimestampType,
2930
DuckDBTinyIntType,
3031
DuckDBType,
@@ -36,6 +37,7 @@ import {
3637
DuckDBUUIDType,
3738
DuckDBUnionType,
3839
DuckDBVarCharType,
40+
DuckDBVarIntType,
3941
} from './DuckDBType';
4042
import { DuckDBTypeId } from './DuckDBTypeId';
4143

@@ -122,6 +124,26 @@ function getBuffer(dataView: DataView, offset: number): Buffer | null {
122124
return Buffer.from(stringBytes);
123125
}
124126

127+
function getVarIntFromBytes(bytes: Uint8Array): bigint {
128+
const firstByte = bytes[0];
129+
const positive = (firstByte & 0x80) > 0;
130+
const uint64Mask = positive ? 0n : 0xffffffffffffffffn;
131+
const uint8Mask = positive ? 0 : 0xff;
132+
const dv = new DataView(bytes.buffer, bytes.byteOffset + 3, bytes.byteLength - 3);
133+
const lastUint64Offset = dv.byteLength - 8;
134+
let offset = 0;
135+
let result = 0n;
136+
while (offset <= lastUint64Offset) {
137+
result = (result << 64n) | (dv.getBigUint64(offset) ^ uint64Mask);
138+
offset += 8;
139+
}
140+
while (offset < dv.byteLength) {
141+
result = (result << 8n) | BigInt(dv.getUint8(offset) ^ uint8Mask);
142+
offset += 1;
143+
}
144+
return positive ? result : -result;
145+
}
146+
125147
function getBoolean1(dataView: DataView, offset: number): boolean {
126148
return getUInt8(dataView, offset) !== 0
127149
}
@@ -360,13 +382,13 @@ export abstract class DuckDBVector<T> {
360382
case DuckDBTypeId.TIME_TZ:
361383
return DuckDBTimeTZVector.fromRawVector(vector, itemCount);
362384
case DuckDBTypeId.TIMESTAMP_TZ:
363-
return DuckDBTimestampVector.fromRawVector(vector, itemCount);
385+
return DuckDBTimestampTZVector.fromRawVector(vector, itemCount);
364386
case DuckDBTypeId.ANY:
365-
throw new Error(`Vector not implemented for ANY type`);
387+
throw new Error(`Invalid vector type: ANY`);
366388
case DuckDBTypeId.VARINT:
367-
return DuckDBBlobVector.fromRawVector(vector, itemCount); // TODO: VARINT
389+
return DuckDBVarIntVector.fromRawVector(vector, itemCount);
368390
case DuckDBTypeId.SQLNULL:
369-
throw new Error(`Vector not implemented for SQLNULL type`);
391+
throw new Error(`Invalid vector type: SQLNULL`);
370392
default:
371393
throw new Error(`Invalid type id: ${vectorType.typeId}`);
372394
}
@@ -1774,3 +1796,69 @@ export class DuckDBTimeTZVector extends DuckDBVector<DuckDBTimeTZValue> {
17741796
return new DuckDBTimeTZVector(this.items.slice(offset, offset + length), this.validity.slice(offset));
17751797
}
17761798
}
1799+
1800+
export class DuckDBTimestampTZVector extends DuckDBVector<bigint> {
1801+
private readonly items: BigInt64Array;
1802+
private readonly validity: DuckDBValidity;
1803+
constructor(items: BigInt64Array, validity: DuckDBValidity) {
1804+
super();
1805+
this.items = items;
1806+
this.validity = validity;
1807+
}
1808+
static fromRawVector(vector: duckdb.Vector, itemCount: number): DuckDBTimestampTZVector {
1809+
const data = vectorData(vector, itemCount * BigInt64Array.BYTES_PER_ELEMENT);
1810+
const items = new BigInt64Array(data.buffer, data.byteOffset, itemCount);
1811+
const validity = DuckDBValidity.fromVector(vector, itemCount);
1812+
return new DuckDBTimestampTZVector(items, validity);
1813+
}
1814+
public override get type(): DuckDBTimestampType {
1815+
return DuckDBTimestampTZType.instance;
1816+
}
1817+
public override get itemCount(): number {
1818+
return this.items.length;
1819+
}
1820+
public override getItem(itemIndex: number): bigint | null { // microseconds
1821+
return this.validity.itemValid(itemIndex) ? this.items[itemIndex] : null;
1822+
}
1823+
public override slice(offset: number, length: number): DuckDBTimestampTZVector {
1824+
return new DuckDBTimestampTZVector(this.items.slice(offset, offset + length), this.validity.slice(offset));
1825+
}
1826+
}
1827+
1828+
export class DuckDBVarIntVector extends DuckDBVector<bigint> {
1829+
private readonly dataView: DataView;
1830+
private readonly validity: DuckDBValidity;
1831+
private readonly _itemCount: number;
1832+
constructor(dataView: DataView, validity: DuckDBValidity, itemCount: number) {
1833+
super();
1834+
this.dataView = dataView;
1835+
this.validity = validity;
1836+
this._itemCount = itemCount;
1837+
}
1838+
static fromRawVector(vector: duckdb.Vector, itemCount: number): DuckDBVarIntVector {
1839+
const data = vectorData(vector, itemCount * 16);
1840+
const dataView = new DataView(data.buffer, data.byteOffset, data.byteLength);
1841+
const validity = DuckDBValidity.fromVector(vector, itemCount);
1842+
return new DuckDBVarIntVector(dataView, validity, itemCount);
1843+
}
1844+
public override get type(): DuckDBVarIntType {
1845+
return DuckDBVarIntType.instance;
1846+
}
1847+
public override get itemCount(): number {
1848+
return this._itemCount;
1849+
}
1850+
public override getItem(itemIndex: number): bigint | null {
1851+
if (!this.validity.itemValid(itemIndex)) {
1852+
return null;
1853+
}
1854+
const bytes = getStringBytes(this.dataView, itemIndex * 16);
1855+
return bytes ? getVarIntFromBytes(bytes) : null;
1856+
}
1857+
public override slice(offset: number, length: number): DuckDBVarIntVector {
1858+
return new DuckDBVarIntVector(
1859+
new DataView(this.dataView.buffer, this.dataView.byteOffset + offset * 16, length * 16),
1860+
this.validity.slice(offset),
1861+
length,
1862+
);
1863+
}
1864+
}

‎api/test/api.test.ts

Lines changed: 8 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,7 @@ import {
6262
DuckDBTimestampSecondsType,
6363
DuckDBTimestampSecondsVector,
6464
DuckDBTimestampTZType,
65+
DuckDBTimestampTZVector,
6566
DuckDBTimestampType,
6667
DuckDBTimestampVector,
6768
DuckDBTinyIntType,
@@ -85,6 +86,7 @@ import {
8586
DuckDBVarCharType,
8687
DuckDBVarCharVector,
8788
DuckDBVarIntType,
89+
DuckDBVarIntVector,
8890
DuckDBVector,
8991
configurationOptionDescriptions,
9092
version
@@ -145,7 +147,7 @@ const MinTimeTZOffset = -MaxTimeTZOffset;
145147
const MinTimeTZ = new DuckDBTimeTZValue(MinTimeTZMicroseconds, MaxTimeTZOffset);
146148
const MaxTimeTZ = new DuckDBTimeTZValue(MaxTimeTZMicroseconds, MinTimeTZOffset);
147149
const MinTS_S = BigInt(-9223372022400); // from test_all_types() select epoch(timestamp_s)::bigint;
148-
const MaxTS_S = BigInt(9223372036854);
150+
const MaxTS_S = BigInt( 9223372036854);
149151
const MinTS_MS = MinTS_S * BI_1000;
150152
const MaxTS_MS = (MaxInt64 - BI_1) / BI_1000;
151153
const MinTS_US = MinTS_MS * BI_1000;
@@ -154,31 +156,13 @@ const TS_US_Inf = MaxInt64;
154156
const MinTS_NS = -9223286400000000000n;
155157
const MaxTS_NS = MaxInt64 - BI_1;
156158
const MinFloat32 = Math.fround(-3.4028235e+38);
157-
const MaxFloat32 = Math.fround(3.4028235e+38);
159+
const MaxFloat32 = Math.fround( 3.4028235e+38);
158160
const MinFloat64 = -Number.MAX_VALUE;
159161
const MaxFloat64 = Number.MAX_VALUE;
160162
const MinUUID = MinInt128;
161163
const MaxUUID = MaxInt128;
162-
const MinVarInt = new Uint8Array([0x7F, 0xFF, 0x7F,
163-
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x07, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
164-
0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
165-
0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
166-
0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
167-
0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
168-
0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
169-
0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
170-
0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
171-
]);
172-
const MaxVarInt = new Uint8Array([0x80, 0x00, 0x80,
173-
0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xF8, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
174-
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
175-
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
176-
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
177-
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
178-
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
179-
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
180-
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
181-
]);
164+
const MinVarInt = -179769313486231570814527423731704356798070567525844996598917476803157260780028538760589558632766878171540458953514382464234321326889464182768467546703537516986049910576551282076245490090389328944075868508455133942304583236903222948165808559332123348274797826204144723168738177180919299881250404026184124858368n
165+
const MaxVarInt = 179769313486231570814527423731704356798070567525844996598917476803157260780028538760589558632766878171540458953514382464234321326889464182768467546703537516986049910576551282076245490090389328944075868508455133942304583236903222948165808559332123348274797826204144723168738177180919299881250404026184124858368n;
182166

183167
async function sleep(ms: number): Promise<void> {
184168
return new Promise((resolve) => {
@@ -544,15 +528,15 @@ describe('api', () => {
544528
assertValues(chunk, 8, DuckDBUSmallIntVector, [MinUInt16, MaxUInt16, null]);
545529
assertValues(chunk, 9, DuckDBUIntegerVector, [MinUInt32, MaxUInt32, null]);
546530
assertValues(chunk, 10, DuckDBUBigIntVector, [MinUInt64, MaxUInt64, null]);
547-
assertValues(chunk, 11, DuckDBBlobVector, [MinVarInt, MaxVarInt, null]);
531+
assertValues(chunk, 11, DuckDBVarIntVector, [MinVarInt, MaxVarInt, null]);
548532
assertValues(chunk, 12, DuckDBDateVector, [MinDate, MaxDate, null]);
549533
assertValues(chunk, 13, DuckDBTimeVector, [MinTime, MaxTime, null]);
550534
assertValues(chunk, 14, DuckDBTimestampVector, [MinTS_US, MaxTS_US, null]);
551535
assertValues(chunk, 15, DuckDBTimestampSecondsVector, [MinTS_S, MaxTS_S, null]);
552536
assertValues(chunk, 16, DuckDBTimestampMillisecondsVector, [MinTS_MS, MaxTS_MS, null]);
553537
assertValues(chunk, 17, DuckDBTimestampNanosecondsVector, [MinTS_NS, MaxTS_NS, null]);
554538
assertValues(chunk, 18, DuckDBTimeTZVector, [MinTimeTZ, MaxTimeTZ, null]);
555-
assertValues(chunk, 19, DuckDBTimestampVector, [MinTS_US, MaxTS_US, null]);
539+
assertValues(chunk, 19, DuckDBTimestampTZVector, [MinTS_US, MaxTS_US, null]);
556540
assertValues(chunk, 20, DuckDBFloatVector, [MinFloat32, MaxFloat32, null]);
557541
assertValues(chunk, 21, DuckDBDoubleVector, [MinFloat64, MaxFloat64, null]);
558542
assertValues(chunk, 22, DuckDBDecimal2Vector, [

0 commit comments

Comments
 (0)