@leonardovida-md/drizzle-neo-duckdb 1.0.1 → 1.0.3

package/src/olap.ts

@@ -0,0 +1,189 @@
+import { is } from 'drizzle-orm/entity';
+import { sql, Subquery, type SQLWrapper } from 'drizzle-orm';
+import type { AnyPgColumn, PgTable } from 'drizzle-orm/pg-core';
+import type { PgViewBase } from 'drizzle-orm/pg-core/view-base';
+import type { SelectedFields } from 'drizzle-orm/pg-core/query-builders';
+import { SQL } from 'drizzle-orm/sql/sql';
+import { Column, getTableName } from 'drizzle-orm';
+import type { DuckDBDatabase } from './driver.ts';
+
+export const countN = (expr: SQLWrapper = sql`*`) =>
+  sql<number>`count(${expr})`.mapWith(Number);
+
+export const sumN = (expr: SQLWrapper) =>
+  sql<number>`sum(${expr})`.mapWith(Number);
+
+export const avgN = (expr: SQLWrapper) =>
+  sql<number>`avg(${expr})`.mapWith(Number);
+
+export const sumDistinctN = (expr: SQLWrapper) =>
+  sql<number>`sum(distinct ${expr})`.mapWith(Number);
+
+export const percentileCont = (p: number, expr: SQLWrapper) =>
+  sql<number>`percentile_cont(${p}) within group (order by ${expr})`.mapWith(
+    Number
+  );
+
+export const median = (expr: SQLWrapper) => percentileCont(0.5, expr);
+
+export const anyValue = <T = unknown>(expr: SQLWrapper) =>
+  sql<T>`any_value(${expr})`;
+
+type PartitionOrder =
+  | {
+      partitionBy?: SQLWrapper | SQLWrapper[];
+      orderBy?: SQLWrapper | SQLWrapper[];
+    }
+  | undefined;
+
+function normalizeArray<T>(value?: T | T[]): T[] {
+  if (!value) return [];
+  return Array.isArray(value) ? value : [value];
+}
+
+function overClause(options?: PartitionOrder) {
+  const partitions = normalizeArray(options?.partitionBy);
+  const orders = normalizeArray(options?.orderBy);
+
+  const chunks: SQLWrapper[] = [];
+
+  if (partitions.length > 0) {
+    chunks.push(sql`partition by ${sql.join(partitions, sql`, `)}`);
+  }
+
+  if (orders.length > 0) {
+    chunks.push(sql`order by ${sql.join(orders, sql`, `)}`);
+  }
+
+  if (chunks.length === 0) {
+    return sql``;
+  }
+
+  return sql`over (${sql.join(chunks, sql` `)})`;
+}
+
+export const rowNumber = (options?: PartitionOrder) =>
+  sql<number>`row_number() ${overClause(options)}`.mapWith(Number);
+
+export const rank = (options?: PartitionOrder) =>
+  sql<number>`rank() ${overClause(options)}`.mapWith(Number);
+
+export const denseRank = (options?: PartitionOrder) =>
+  sql<number>`dense_rank() ${overClause(options)}`.mapWith(Number);
+
+export const lag = <T = unknown>(
+  expr: SQLWrapper,
+  offset = 1,
+  defaultValue?: SQLWrapper,
+  options?: PartitionOrder
+) =>
+  defaultValue
+    ? sql<T>`lag(${expr}, ${offset}, ${defaultValue}) ${overClause(options)}`
+    : sql<T>`lag(${expr}, ${offset}) ${overClause(options)}`;
+
+export const lead = <T = unknown>(
+  expr: SQLWrapper,
+  offset = 1,
+  defaultValue?: SQLWrapper,
+  options?: PartitionOrder
+) =>
+  defaultValue
+    ? sql<T>`lead(${expr}, ${offset}, ${defaultValue}) ${overClause(options)}`
+    : sql<T>`lead(${expr}, ${offset}) ${overClause(options)}`;
+
+type ValueExpr = SQL | SQL.Aliased | AnyPgColumn;
+type GroupKey = ValueExpr;
+type MeasureMap = Record<string, ValueExpr>;
+type NonAggMap = Record<string, ValueExpr>;
+
+function keyAlias(key: SQLWrapper, fallback: string): string {
+  if (is(key, SQL.Aliased)) {
+    return key.fieldAlias ?? fallback;
+  }
+  if (is(key, Column)) {
+    return `${getTableName(key.table)}.${key.name}`;
+  }
+  return fallback;
+}
+
+export class OlapBuilder {
+  private source?: PgTable | Subquery | PgViewBase | SQL;
+  private keys: GroupKey[] = [];
+  private measureMap: MeasureMap = {};
+  private nonAggregates: NonAggMap = {};
+  private wrapNonAggWithAnyValue = false;
+  private orderByClauses: ValueExpr[] = [];
+
+  constructor(private db: DuckDBDatabase) {}
+
+  from(source: PgTable | Subquery | PgViewBase | SQL): this {
+    this.source = source;
+    return this;
+  }
+
+  groupBy(keys: GroupKey[]): this {
+    this.keys = keys;
+    return this;
+  }
+
+  measures(measures: MeasureMap): this {
+    this.measureMap = measures;
+    return this;
+  }
+
+  selectNonAggregates(
+    fields: NonAggMap,
+    options: { anyValue?: boolean } = {}
+  ): this {
+    this.nonAggregates = fields;
+    this.wrapNonAggWithAnyValue = options.anyValue ?? false;
+    return this;
+  }
+
+  orderBy(...clauses: ValueExpr[]): this {
+    this.orderByClauses = clauses;
+    return this;
+  }
+
+  build() {
+    if (!this.source) {
+      throw new Error('olap: .from() is required');
+    }
+    if (this.keys.length === 0) {
+      throw new Error('olap: .groupBy() is required');
+    }
+    if (Object.keys(this.measureMap).length === 0) {
+      throw new Error('olap: .measures() is required');
+    }
+
+    const selection: Record<string, ValueExpr> = {};
+
+    this.keys.forEach((key, idx) => {
+      const alias = keyAlias(key, `key_${idx}`);
+      selection[alias] = key;
+    });
+
+    Object.entries(this.nonAggregates).forEach(([alias, expr]) => {
+      selection[alias] = this.wrapNonAggWithAnyValue ? anyValue(expr) : expr;
+    });
+
+    Object.assign(selection, this.measureMap);
+
+    let query: any = this.db
+      .select(selection as SelectedFields)
+      .from(this.source!)
+      .groupBy(...this.keys);
+
+    if (this.orderByClauses.length > 0) {
+      query = query.orderBy(...this.orderByClauses);
+    }
+
+    return query;
+  }
+
+  run() {
+    return this.build();
+  }
+}
+
+export const olap = (db: DuckDBDatabase) => new OlapBuilder(db);

package/src/select-builder.ts

@@ -6,11 +6,7 @@ import {
   type SelectedFields,
   type TableLikeHasEmptySelection,
 } from 'drizzle-orm/pg-core/query-builders';
-import {
-  PgColumn,
-  PgTable,
-  type PgSession,
-} from 'drizzle-orm/pg-core';
+import { PgColumn, PgTable, type PgSession } from 'drizzle-orm/pg-core';
 import { Subquery, ViewBaseConfig, type SQLWrapper } from 'drizzle-orm';
 import { PgViewBase } from 'drizzle-orm/pg-core/view-base';
 import type {
@@ -25,7 +21,7 @@ import { getTableColumns, type DrizzleTypeError } from 'drizzle-orm/utils';
 interface PgViewBaseInternal<
   TName extends string = string,
   TExisting extends boolean = boolean,
-  TSelectedFields extends ColumnsSelection = ColumnsSelection
+  TSelectedFields extends ColumnsSelection = ColumnsSelection,
 > extends PgViewBase<TName, TExisting, TSelectedFields> {
   [ViewBaseConfig]?: {
     selectedFields: SelectedFields;
@@ -34,7 +30,7 @@ interface PgViewBaseInternal<
 
 export class DuckDBSelectBuilder<
   TSelection extends SelectedFields | undefined,
-  TBuilderMode extends 'db' | 'qb' = 'db'
+  TBuilderMode extends 'db' | 'qb' = 'db',
 > extends PgSelectBuilder<TSelection, TBuilderMode> {
   private _fields: TSelection;
   private _session: PgSession | undefined;

package/src/session.ts

@@ -20,12 +20,18 @@ import { mapResultRow } from './sql/result-mapper.ts';
 import type { DuckDBDialect } from './dialect.ts';
 import { TransactionRollbackError } from 'drizzle-orm/errors';
 import type { DuckDBClientLike, RowData } from './client.ts';
-import { executeOnClient, prepareParams } from './client.ts';
+import {
+  executeArrowOnClient,
+  executeInBatches,
+  executeOnClient,
+  prepareParams,
+  type ExecuteInBatchesOptions,
+} from './client.ts';
 
 export type { DuckDBClientLike, RowData } from './client.ts';
 
 export class DuckDBPreparedQuery<
-  T extends PreparedQueryConfig
+  T extends PreparedQueryConfig,
 > extends PgPreparedQuery<T> {
   static readonly [entityKind]: string = 'DuckDBPreparedQuery';
 
@@ -37,7 +43,9 @@ export class DuckDBPreparedQuery<
     private logger: Logger,
     private fields: SelectedFieldsOrdered | undefined,
     private _isResponseInArrayMode: boolean,
-    private customResultMapper: ((rows: unknown[][]) => T['execute']) | undefined,
+    private customResultMapper:
+      | ((rows: unknown[][]) => T['execute'])
+      | undefined,
     private rewriteArrays: boolean,
     private rejectStringArrayLiterals: boolean,
     private warnOnStringArrayLiteral?: (sql: string) => void
@@ -71,17 +79,10 @@ export class DuckDBPreparedQuery<
 
     this.logger.logQuery(rewrittenQuery, params);
 
-    const {
-      fields,
-      joinsNotNullableMap,
-      customResultMapper,
-    } = this as typeof this & { joinsNotNullableMap?: Record<string, boolean> };
+    const { fields, joinsNotNullableMap, customResultMapper } =
+      this as typeof this & { joinsNotNullableMap?: Record<string, boolean> };
 
-    const rows = await executeOnClient(
-      this.client,
-      rewrittenQuery,
-      params
-    );
+    const rows = await executeOnClient(this.client, rewrittenQuery, params);
 
     if (rows.length === 0 || !fields) {
       return rows as T['execute'];
@@ -115,7 +116,7 @@ export interface DuckDBSessionOptions {
 
 export class DuckDBSession<
   TFullSchema extends Record<string, unknown> = Record<string, never>,
-  TSchema extends TablesRelationalConfig = Record<string, never>
+  TSchema extends TablesRelationalConfig = Record<string, never>,
 > extends PgSession<DuckDBQueryResultHKT, TFullSchema, TSchema> {
   static readonly [entityKind]: string = 'DuckDBSession';
 
@@ -199,11 +200,67 @@ export class DuckDBSession<
       []
     );
   };
+
+  executeBatches<T extends RowData = RowData>(
+    query: SQL,
+    options: ExecuteInBatchesOptions = {}
+  ): AsyncGenerator<GenericRowData<T>[], void, void> {
+    const builtQuery = this.dialect.sqlToQuery(query);
+    const params = prepareParams(builtQuery.params, {
+      rejectStringArrayLiterals: this.rejectStringArrayLiterals,
+      warnOnStringArrayLiteral: this.rejectStringArrayLiterals
+        ? undefined
+        : () => this.warnOnStringArrayLiteral(builtQuery.sql),
+    });
+    const rewrittenQuery = this.rewriteArrays
+      ? adaptArrayOperators(builtQuery.sql)
+      : builtQuery.sql;
+
+    if (this.rewriteArrays && rewrittenQuery !== builtQuery.sql) {
+      this.logger.logQuery(
+        `[duckdb] original query before array rewrite: ${builtQuery.sql}`,
+        params
+      );
+    }
+
+    this.logger.logQuery(rewrittenQuery, params);
+
+    return executeInBatches(
+      this.client,
+      rewrittenQuery,
+      params,
+      options
+    ) as AsyncGenerator<GenericRowData<T>[], void, void>;
+  }
+
+  async executeArrow(query: SQL): Promise<unknown> {
+    const builtQuery = this.dialect.sqlToQuery(query);
+    const params = prepareParams(builtQuery.params, {
+      rejectStringArrayLiterals: this.rejectStringArrayLiterals,
+      warnOnStringArrayLiteral: this.rejectStringArrayLiterals
+        ? undefined
+        : () => this.warnOnStringArrayLiteral(builtQuery.sql),
+    });
+    const rewrittenQuery = this.rewriteArrays
+      ? adaptArrayOperators(builtQuery.sql)
+      : builtQuery.sql;
+
+    if (this.rewriteArrays && rewrittenQuery !== builtQuery.sql) {
+      this.logger.logQuery(
+        `[duckdb] original query before array rewrite: ${builtQuery.sql}`,
+        params
+      );
+    }
+
+    this.logger.logQuery(rewrittenQuery, params);
+
+    return executeArrowOnClient(this.client, rewrittenQuery, params);
+  }
 }
 
 type PgTransactionInternals<
   TFullSchema extends Record<string, unknown> = Record<string, never>,
-  TSchema extends TablesRelationalConfig = Record<string, never>
+  TSchema extends TablesRelationalConfig = Record<string, never>,
 > = {
   dialect: DuckDBDialect;
   session: DuckDBSession<TFullSchema, TSchema>;
@@ -211,13 +268,13 @@ type PgTransactionInternals<
 
 type DuckDBTransactionWithInternals<
   TFullSchema extends Record<string, unknown> = Record<string, never>,
-  TSchema extends TablesRelationalConfig = Record<string, never>
+  TSchema extends TablesRelationalConfig = Record<string, never>,
 > = PgTransactionInternals<TFullSchema, TSchema> &
   DuckDBTransaction<TFullSchema, TSchema>;
 
 export class DuckDBTransaction<
   TFullSchema extends Record<string, unknown>,
-  TSchema extends TablesRelationalConfig
+  TSchema extends TablesRelationalConfig,
 > extends PgTransaction<DuckDBQueryResultHKT, TFullSchema, TSchema> {
   static readonly [entityKind]: string = 'DuckDBTransaction';
 
@@ -246,6 +303,19 @@ export class DuckDBTransaction<
     );
   }
 
+  executeBatches<T extends RowData = RowData>(
+    query: SQL,
+    options: ExecuteInBatchesOptions = {}
+  ): AsyncGenerator<GenericRowData<T>[], void, void> {
+    type Tx = DuckDBTransactionWithInternals<TFullSchema, TSchema>;
+    return (this as unknown as Tx).session.executeBatches<T>(query, options);
+  }
+
+  executeArrow(query: SQL): Promise<unknown> {
+    type Tx = DuckDBTransactionWithInternals<TFullSchema, TSchema>;
+    return (this as unknown as Tx).session.executeArrow(query);
+  }
+
   override async transaction<T>(
     transaction: (tx: DuckDBTransaction<TFullSchema, TSchema>) => Promise<T>
   ): Promise<T> {
@@ -268,7 +338,6 @@ export type GenericTableData<T = RowData> = T[];
 const arrayLiteralWarning =
   'Received a stringified Postgres-style array literal. Use duckDbList()/duckDbArray() or pass native arrays instead. You can also set rejectStringArrayLiterals=true to throw.';
 
-
 export interface DuckDBQueryResultHKT extends PgQueryResultHKT {
   type: GenericTableData<Assume<this['row'], RowData>>;
 }

package/src/sql/query-rewriters.ts

@@ -85,10 +85,7 @@ export function adaptArrayOperators(query: string): string {
     let idx = rewritten.indexOf(token);
     while (idx !== -1) {
       const [leftStart, leftExpr] = walkLeft(rewritten, idx - 1);
-      const [rightEnd, rightExpr] = walkRight(
-        rewritten,
-        idx + token.length
-      );
+      const [rightEnd, rightExpr] = walkRight(rewritten, idx + token.length);
 
       const left = leftExpr.trim();
       const right = rightExpr.trim();
@@ -98,9 +95,7 @@ export function adaptArrayOperators(query: string): string {
       })`;
 
       rewritten =
-        rewritten.slice(0, leftStart) +
-        replacement +
-        rewritten.slice(rightEnd);
+        rewritten.slice(0, leftStart) + replacement + rewritten.slice(rightEnd);
 
       idx = rewritten.indexOf(token, leftStart + replacement.length);
     }
@@ -136,7 +131,9 @@ export function queryAdapter(query: string): string {
       if (noTableProp) {
         const [, column, alias] = noTableProp;
         const asAlias = ` as '${column}'`;
-        return alias ? trimmedField.replace(alias, asAlias) : `${trimmedField}${asAlias}`;
+        return alias
+          ? trimmedField.replace(alias, asAlias)
+          : `${trimmedField}${asAlias}`;
       }
 
       return trimmedField;

package/src/sql/result-mapper.ts

@@ -88,7 +88,10 @@ function normalizeTimestampString(
   return value;
 }
 
-function normalizeTimestamp(value: unknown, withTimezone: boolean): Date | unknown {
+function normalizeTimestamp(
+  value: unknown,
+  withTimezone: boolean
+): Date | unknown {
   if (value instanceof Date) {
     return value;
   }
@@ -96,8 +99,7 @@ function normalizeTimestamp(value: unknown, withTimezone: boolean): Date | unkno
     const hasOffset =
       value.endsWith('Z') || /[+-]\d{2}:?\d{2}$/.test(value.trim());
     const spaced = value.replace(' ', 'T');
-    const normalized =
-      withTimezone || hasOffset ? spaced : `${spaced}+00`;
+    const normalized = withTimezone || hasOffset ? spaced : `${spaced}+00`;
     return new Date(normalized);
   }
   return value;
@@ -177,9 +179,7 @@ function mapDriverValue(
     if (normalized instanceof Date) {
       return normalized;
     }
-    return decoder.mapFromDriverValue(
-      toDecoderInput(decoder, normalized)
-    );
+    return decoder.mapFromDriverValue(toDecoderInput(decoder, normalized));
   }
 
   if (is(decoder, PgDateString)) {

package/src/sql/selection.ts

@@ -1,10 +1,4 @@
-import {
-  Column,
-  SQL,
-  getTableName,
-  is,
-  sql,
-} from 'drizzle-orm';
+import { Column, SQL, getTableName, is, sql } from 'drizzle-orm';
 import type { SelectedFields } from 'drizzle-orm/pg-core';
 
 function mapEntries(
@@ -38,8 +32,7 @@ function mapEntries(
         }
 
         if (is(value, SQL) || is(value, Column)) {
-          const aliased =
-            is(value, SQL) ? value : sql`${value}`.mapWith(value);
+          const aliased = is(value, SQL) ? value : sql`${value}`.mapWith(value);
           return [key, aliased.as(qualified)];
         }