Line data Source code
1 : import type * as arrow from 'apache-arrow';
2 :
3 : import { Logger } from '../platform/logger/logger.js';
4 : import type {
5 : EmbeddedClientContext,
6 : EmbeddedComputeDatabase,
7 : EmbeddedTableImportConnection,
8 : } from '../platform/database/embedded_database.js';
9 : import { OrderByConstraint, SQLFrame } from './sql/sqlframe_builder.js';
10 :
11 1 : const LOG_CTX = "data_frame";
12 :
13 1 : let nextTableId = 0;
14 :
15 : /// The options of a read of the rows of a data frame
16 : export interface DataFrameReadOptions {
17 : /// The order of the rows. A table does not keep the order of the query that created it, the order of a read
18 : /// without constraints is arbitrary.
19 : orderBy?: OrderByConstraint[];
20 : /// The columns to read, all columns by default
21 : columns?: string[];
22 : /// Cancels the read
23 : abort?: AbortSignal;
24 : }
25 :
26 1 : export function generateTableName(prefix: string = "__df"): string {
27 1 : return `${prefix}_${nextTableId++}`;
28 1 : }
29 :
30 1 : export class DataFrame {
31 1 : readonly database: EmbeddedComputeDatabase;
32 1 : readonly tableName: string;
33 : /// The client context of the queries of the data frame, the one of the query that created it
34 1 : readonly clientContext: EmbeddedClientContext | null;
35 :
36 1 : constructor(database: EmbeddedComputeDatabase, tableName: string, clientContext: EmbeddedClientContext | null = null) {
37 1 : this.database = database;
38 1 : this.tableName = tableName;
39 1 : this.clientContext = clientContext;
40 1 : }
41 :
42 1 : static async withConnection<T>(
43 1 : database: EmbeddedComputeDatabase,
44 1 : fn: (conn: EmbeddedTableImportConnection) => Promise<T>,
45 1 : clientContext?: EmbeddedClientContext | null,
46 1 : ): Promise<T> {
47 1 : const conn = await database.connect({ pooled: true, clientContext: clientContext ?? undefined });
48 1 : try {
49 1 : return await fn(conn);
50 1 : } finally {
51 1 : await conn.close();
52 1 : }
53 1 : }
54 :
55 1 : async withConnection<T>(fn: (conn: EmbeddedTableImportConnection) => Promise<T>): Promise<T> {
56 1 : return await DataFrame.withConnection(this.database, fn, this.clientContext);
57 1 : }
58 :
59 1 : static async fromArrowTable(
60 1 : database: EmbeddedComputeDatabase,
61 1 : table: arrow.Table,
62 1 : tableName: string,
63 1 : clientContext: EmbeddedClientContext | null = null,
64 1 : ): Promise<DataFrame> {
65 1 : await DataFrame.withConnection(database, async conn => {
66 1 : await conn.insertArrowTable(table, { name: tableName, create: true });
67 1 : }, clientContext);
68 1 : return new DataFrame(database, tableName, clientContext);
69 1 : }
70 :
71 1 : static async fromSQL(
72 1 : database: EmbeddedComputeDatabase,
73 1 : sql: string,
74 1 : tableName: string,
75 1 : abort?: AbortSignal,
76 1 : clientContext?: EmbeddedClientContext | null,
77 1 : ): Promise<DataFrame> {
78 1 : await DataFrame.withConnection(database, async conn => {
79 1 : await conn.createTableAs(tableName, sql, abort);
80 1 : }, clientContext);
81 1 : return new DataFrame(database, tableName, clientContext);
82 1 : }
83 :
84 1 : async readTable(options: DataFrameReadOptions = {}): Promise<arrow.Table> {
85 1 : let frame = SQLFrame.from(this.tableName);
86 1 : if (options.orderBy != null) frame = frame.orderBy(options.orderBy);
87 1 : if (options.columns != null) frame = frame.project(options.columns);
88 1 : const sql = frame.toSQL();
89 1 : return await this.withConnection(async conn => await conn.query(sql, options.abort));
90 1 : }
91 :
92 1 : async destroy(): Promise<void> {
93 1 : await this.withConnection(async conn => await conn.execute(`DROP TABLE IF EXISTS "${this.tableName}"`));
94 1 : }
95 1 : }
96 :
97 1 : export class DataFrameRegistry {
98 1 : logger: Logger;
99 1 : registeredDataFrames: Map<DataFrame, number> = new Map();
100 :
101 1 : constructor(logger: Logger) {
102 1 : this.logger = logger;
103 1 : }
104 :
105 1 : public getRegisteredDataFrames() {
106 1 : return this.registeredDataFrames;
107 1 : }
108 :
109 1 : acquire(dataFrame: DataFrame | null | undefined, times: number = 1) {
110 1 : if (dataFrame == undefined || dataFrame == null) {
111 1 : return;
112 : }
113 1 : if (this.registeredDataFrames.has(dataFrame)) {
114 1 : this.registeredDataFrames.set(dataFrame, this.registeredDataFrames.get(dataFrame)! + times);
115 : } else {
116 1 : this.registeredDataFrames.set(dataFrame, times);
117 : }
118 1 : }
119 :
120 1 : release(dataFrame?: DataFrame | null) {
121 1 : if (dataFrame == undefined || dataFrame == null) {
122 1 : return;
123 : }
124 1 : if (this.registeredDataFrames.has(dataFrame)) {
125 1 : const count = this.registeredDataFrames.get(dataFrame)! - 1;
126 1 : if (count <= 0) {
127 1 : this.registeredDataFrames.delete(dataFrame);
128 1 : void dataFrame.destroy().catch(error => {
129 1 : this.logger.warn("Failed to destroy released data frame", {
130 1 : "tableName": dataFrame.tableName,
131 1 : "error": error instanceof Error ? error.message : String(error),
132 1 : }, LOG_CTX);
133 1 : });
134 1 : } else {
135 1 : this.registeredDataFrames.set(dataFrame, count);
136 : }
137 1 : } else {
138 1 : this.logger.warn("Attempted to release unknown data frame", {
139 1 : "tableName": dataFrame.tableName
140 1 : }, LOG_CTX);
141 : }
142 1 : }
143 :
144 1 : releaseMany(dataFrames: (DataFrame | null | undefined)[]) {
145 1 : for (const dataFrame of dataFrames) {
146 1 : this.release(dataFrame);
147 : }
148 1 : }
149 1 : }
|
