pipework 0.252.0 → 0.253.0

package/src/data/query/pg-core/columns/vector_extension/sparsevec.ts

@@ -0,0 +1,67 @@
+import type { ColumnBuilderBaseConfig, ColumnBuilderRuntimeConfig, MakeColumnConfig } from '../../../column-builder.js';
+import type { ColumnBaseConfig } from '../../../column.js';
+import { entityKind } from '../../../entity.js';
+import type { AnyPgTable } from '../../table.js';
+import { getColumnNameAndConfig } from '../../../utils.js';
+import { PgColumn, PgColumnBuilder } from '../common.js';
+
+export type PgSparseVectorBuilderInitial<TName extends string> = PgSparseVectorBuilder<{
+	name: TName;
+	dataType: 'string';
+	columnType: 'PgSparseVector';
+	data: string;
+	driverParam: string;
+	enumValues: undefined;
+}>;
+
+export class PgSparseVectorBuilder<T extends ColumnBuilderBaseConfig<'string', 'PgSparseVector'>>
+	extends PgColumnBuilder<
+		T,
+		{ dimensions: number | undefined }
+	>
+{
+	static override readonly [entityKind]: string = 'PgSparseVectorBuilder';
+
+	constructor(name: string, config: PgSparseVectorConfig) {
+		super(name, 'string', 'PgSparseVector');
+		this.config.dimensions = config.dimensions;
+	}
+
+	/** @internal */
+	override build<TTableName extends string>(
+		table: AnyPgTable<{ name: TTableName }>,
+	): PgSparseVector<MakeColumnConfig<T, TTableName>> {
+		return new PgSparseVector<MakeColumnConfig<T, TTableName>>(
+			table,
+			this.config as ColumnBuilderRuntimeConfig<any, any>,
+		);
+	}
+}
+
+export class PgSparseVector<T extends ColumnBaseConfig<'string', 'PgSparseVector'>>
+	extends PgColumn<T, { dimensions: number | undefined }>
+{
+	static override readonly [entityKind]: string = 'PgSparseVector';
+
+	readonly dimensions = this.config.dimensions;
+
+	getSQLType(): string {
+		return `sparsevec(${this.dimensions})`;
+	}
+}
+
+export interface PgSparseVectorConfig {
+	dimensions: number;
+}
+
+export function sparsevec(
+	config: PgSparseVectorConfig,
+): PgSparseVectorBuilderInitial<''>;
+export function sparsevec<TName extends string>(
+	name: TName,
+	config: PgSparseVectorConfig,
+): PgSparseVectorBuilderInitial<TName>;
+export function sparsevec(a: string | PgSparseVectorConfig, b?: PgSparseVectorConfig) {
+	const { name, config } = getColumnNameAndConfig<PgSparseVectorConfig>(a, b);
+	return new PgSparseVectorBuilder(name, config);
+}

package/src/data/query/pg-core/columns/vector_extension/vector.ts

@@ -0,0 +1,80 @@
+import type { ColumnBuilderBaseConfig, ColumnBuilderRuntimeConfig, MakeColumnConfig } from '../../../column-builder.js';
+import type { ColumnBaseConfig } from '../../../column.js';
+import { entityKind } from '../../../entity.js';
+import type { AnyPgTable } from '../../table.js';
+import { getColumnNameAndConfig } from '../../../utils.js';
+import { PgColumn, PgColumnBuilder } from '../common.js';
+
+export type PgVectorBuilderInitial<TName extends string, TDimensions extends number> = PgVectorBuilder<{
+	name: TName;
+	dataType: 'array';
+	columnType: 'PgVector';
+	data: number[];
+	driverParam: string;
+	enumValues: undefined;
+	dimensions: TDimensions;
+}>;
+
+export class PgVectorBuilder<T extends ColumnBuilderBaseConfig<'array', 'PgVector'> & { dimensions: number }>
+	extends PgColumnBuilder<
+		T,
+		{ dimensions: T['dimensions'] },
+		{ dimensions: T['dimensions'] }
+	>
+{
+	static override readonly [entityKind]: string = 'PgVectorBuilder';
+
+	constructor(name: string, config: PgVectorConfig<T['dimensions']>) {
+		super(name, 'array', 'PgVector');
+		this.config.dimensions = config.dimensions;
+	}
+
+	/** @internal */
+	override build<TTableName extends string>(
+		table: AnyPgTable<{ name: TTableName }>,
+	): PgVector<MakeColumnConfig<T, TTableName> & { dimensions: T['dimensions'] }> {
+		return new PgVector<MakeColumnConfig<T, TTableName> & { dimensions: T['dimensions'] }>(
+			table,
+			this.config as ColumnBuilderRuntimeConfig<any, any>,
+		);
+	}
+}
+
+export class PgVector<T extends ColumnBaseConfig<'array', 'PgVector'> & { dimensions: number | undefined }>
+	extends PgColumn<T, { dimensions: T['dimensions'] }, { dimensions: T['dimensions'] }>
+{
+	static override readonly [entityKind]: string = 'PgVector';
+
+	readonly dimensions: T['dimensions'] = this.config.dimensions;
+
+	getSQLType(): string {
+		return `vector(${this.dimensions})`;
+	}
+
+	override mapToDriverValue(value: unknown): unknown {
+		return JSON.stringify(value);
+	}
+
+	override mapFromDriverValue(value: string): unknown {
+		return value
+			.slice(1, -1)
+			.split(',')
+			.map((v) => Number.parseFloat(v));
+	}
+}
+
+export interface PgVectorConfig<TDimensions extends number = number> {
+	dimensions: TDimensions;
+}
+
+export function vector<D extends number>(
+	config: PgVectorConfig<D>,
+): PgVectorBuilderInitial<'', D>;
+export function vector<TName extends string, D extends number>(
+	name: TName,
+	config: PgVectorConfig<D>,
+): PgVectorBuilderInitial<TName, D>;
+export function vector(a: string | PgVectorConfig, b?: PgVectorConfig) {
+	const { name, config } = getColumnNameAndConfig<PgVectorConfig>(a, b);
+	return new PgVectorBuilder(name, config);
+}