Line data Source code
1 : import * as React from 'react';
2 :
3 : import { computeColumnAggregates, computeFilteredColumnAggregates } from './tasks/column_aggregation.js';
4 : import { computeSystemColumns } from './tasks/system_columns.js';
5 : import { computeTableAggregates } from './tasks/table_aggregation.js';
6 : import { filterTable } from './tasks/table_filtering.js';
7 : import { searchData } from './tasks/data_search.js';
8 : import { sortTable } from './tasks/table_ordering.js';
9 :
10 : import { COLUMN_AGGREGATION_SUCCEEDED, ComputationAction, DATA_SEARCH_FAILED, DATA_SEARCH_SUCCEEDED, UNREGISTER_SCHEDULER_TASK, FILTERED_COLUMN_AGGREGATION_SUCCEEDED, SYSTEM_COLUMN_COMPUTATION_SUCCEEDED, TABLE_AGGREGATION_SUCCEEDED, TABLE_FILTERING_SUCCEEDED, TABLE_ORDERING_SUCCEDED, UPDATE_SCHEDULER_TASK } from './computation_state.js';
11 : import { Dispatch } from '../utils/variant.js';
12 : import { LoggableException, Logger, stringifyError } from '../platform/logger/logger.js';
13 : import { createTrace } from '../platform/logger/trace_context.js';
14 : import { TaskStatus, TaskProgress, TaskVariant, COLUMN_AGGREGATION_TASK, DATA_SEARCH_TASK, FILTERED_COLUMN_AGGREGATION_TASK, SYSTEM_COLUMN_COMPUTATION_TASK, TABLE_AGGREGATION_TASK, TABLE_FILTERING_TASK, TABLE_ORDERING_TASK } from "./computation_types.js";
15 : import { useComputationRegistry } from "./computation_registry.js";
16 : import { useLogger } from '../platform/logger/logger_provider.js';
17 : import type { ComputeQueryExecution } from './computation_tasks.js';
18 : import { useCreateComputeQueryExecution } from './compute_query_execution.js';
19 :
20 1 : const LOG_CTX = 'scheduler';
21 :
22 1 : export const COMPUTATION_OPERATIONS = {
23 1 : computeColumnAggregates,
24 1 : computeFilteredColumnAggregates,
25 1 : computeSystemColumns,
26 1 : computeTableAggregates,
27 1 : filterTable,
28 1 : searchData,
29 1 : sortTable,
30 1 : };
31 :
32 : interface SchedulerState {
33 : /// Started tasks
34 : launched: Set<number>
35 : }
36 :
37 1 : export function ComputationScheduler(props: React.PropsWithChildren<{}>) {
38 1 : const logger = useLogger();
39 1 : const [computationState, dispatchComputation] = useComputationRegistry();
40 1 : const createQueryExecution = useCreateComputeQueryExecution();
41 :
42 1 : const schedulerState = React.useRef<SchedulerState>({ launched: new Set() });
43 1 : React.useEffect(() => {
44 1 : const launched = schedulerState.current.launched!;
45 1 : for (const [taskId, task] of Object.entries(computationState.schedulerTasks)) {
46 1 : const id = +taskId;
47 : // Already processed?
48 1 : if (launched.has(id)) {
49 1 : continue;
50 : }
51 : // Launch the task
52 1 : launched.add(id);
53 :
54 : // Process a task asynchronously
55 1 : processTask(task, dispatchComputation, logger, createQueryExecution?.(task));
56 1 : }
57 : // Cleanup tasks that are no longer in included in the background tasks
58 1 : for (const l of launched) {
59 1 : if (computationState.schedulerTasks[l] === undefined) {
60 1 : launched.delete(l);
61 : }
62 : }
63 1 : });
64 :
65 1 : return props.children;
66 1 : }
67 :
68 1 : export async function processTask(
69 1 : task: TaskVariant,
70 1 : dispatchComputation: Dispatch<ComputationAction>,
71 1 : logger: Logger,
72 1 : executeQuery?: ComputeQueryExecution,
73 1 : ) {
74 1 : if (task.taskId === undefined) {
75 1 : logger.warn("Task has no task id", {
76 1 : taskId: task.taskId,
77 1 : taskType: getTaskVariantName(task)
78 1 : }, LOG_CTX);
79 1 : return;
80 1 : }
81 :
82 : // Start a trace for this task
83 1 : const traced = logger.withTrace(createTrace());
84 :
85 1 : let progress: TaskProgress = {
86 1 : status: TaskStatus.TASK_RUNNING,
87 1 : startedAt: new Date(),
88 1 : completedAt: null,
89 1 : failedAt: null,
90 1 : failedWithError: null
91 1 : };
92 1 : try {
93 : // Mark the task as running
94 1 : dispatchComputation({
95 1 : type: UPDATE_SCHEDULER_TASK,
96 1 : value: [task, progress]
97 1 : });
98 :
99 : // Process task
100 1 : switch (task.type) {
101 1 : case TABLE_FILTERING_TASK: {
102 : // Filter the table
103 1 : const filter = await COMPUTATION_OPERATIONS.filterTable(task.value, traced, executeQuery);
104 : // Mark as succeeded
105 1 : dispatchComputation({
106 1 : type: TABLE_FILTERING_SUCCEEDED,
107 1 : value: [task.value.tableId, task.value.tableVersion, filter]
108 1 : });
109 : // Resolve the promise
110 1 : task.result.resolve(filter);
111 1 : break;
112 1 : }
113 1 : case TABLE_ORDERING_TASK: {
114 : // Sort the table
115 1 : const ordered = await COMPUTATION_OPERATIONS.sortTable(task.value, traced, executeQuery);
116 : // Mark as succeeded
117 1 : dispatchComputation({
118 1 : type: TABLE_ORDERING_SUCCEDED,
119 1 : value: [task.value.tableId, ordered]
120 1 : });
121 : // Resolve the promise
122 1 : task.result.resolve(ordered);
123 1 : break;
124 1 : }
125 1 : case DATA_SEARCH_TASK: {
126 1 : const searchTable = await COMPUTATION_OPERATIONS.searchData(task.value, traced, executeQuery);
127 1 : dispatchComputation({
128 1 : type: DATA_SEARCH_SUCCEEDED,
129 1 : value: [task.value.tableId, task.value.inputDataFrame, task.value.requestId, searchTable],
130 1 : });
131 1 : task.result.resolve(searchTable);
132 1 : break;
133 1 : }
134 1 : case TABLE_AGGREGATION_TASK: {
135 : // Aggregate the table
136 1 : const [tableAgg, colEntries] = await COMPUTATION_OPERATIONS.computeTableAggregates(task.value, traced, executeQuery);
137 : // Mark as succeeded
138 1 : dispatchComputation({
139 1 : type: TABLE_AGGREGATION_SUCCEEDED,
140 1 : value: [task.value.tableId, tableAgg]
141 1 : });
142 : // Resolve the promise
143 1 : task.result.resolve([tableAgg, colEntries]);
144 1 : break;
145 1 : }
146 1 : case SYSTEM_COLUMN_COMPUTATION_TASK: {
147 : // Compute the system columns
148 1 : const [table, dataFrame, columnGroups] = await COMPUTATION_OPERATIONS.computeSystemColumns(task.value, traced, executeQuery);
149 : // Mark as succeeded
150 1 : dispatchComputation({
151 1 : type: SYSTEM_COLUMN_COMPUTATION_SUCCEEDED,
152 1 : value: [task.value.tableId, table, dataFrame, columnGroups]
153 1 : });
154 : // Resolve the promise
155 1 : task.result.resolve([table, dataFrame, columnGroups]);
156 1 : break;
157 1 : }
158 1 : case COLUMN_AGGREGATION_TASK:
159 : // Compute column aggregates
160 1 : const columnAgg = await COMPUTATION_OPERATIONS.computeColumnAggregates(task.value, traced, executeQuery);
161 : // Mark as succeeded
162 1 : dispatchComputation({
163 1 : type: COLUMN_AGGREGATION_SUCCEEDED,
164 1 : value: [task.value.tableId, task.value.columnId, columnAgg]
165 1 : });
166 : // Resolve the task
167 1 : task.result.resolve(columnAgg);
168 1 : break;
169 1 : case FILTERED_COLUMN_AGGREGATION_TASK:
170 : // Filtered column aggregates
171 1 : const filteredColumnAgg = await COMPUTATION_OPERATIONS.computeFilteredColumnAggregates(task.value, traced, executeQuery);
172 : // Mark as succeeded
173 1 : dispatchComputation({
174 1 : type: FILTERED_COLUMN_AGGREGATION_SUCCEEDED,
175 1 : value: [task.value.tableId, task.value.columnId, task.value.tableVersion, filteredColumnAgg]
176 1 : });
177 : // Resolve the task
178 1 : task.result.resolve(filteredColumnAgg);
179 : break;
180 1 : }
181 1 : } catch (e: any) {
182 1 : if (task.type === DATA_SEARCH_TASK) {
183 1 : dispatchComputation({
184 1 : type: DATA_SEARCH_FAILED,
185 1 : value: [task.value.tableId, task.value.inputDataFrame, task.value.requestId, stringifyError(e)],
186 1 : });
187 : }
188 1 : let loggable: LoggableException = (e instanceof LoggableException)
189 1 : ? e
190 1 : : new LoggableException("Failed to execute task", { error: stringifyError(e) }, LOG_CTX);
191 :
192 : // Mark the task as failed
193 1 : progress = {
194 1 : ...progress,
195 1 : status: TaskStatus.TASK_FAILED,
196 1 : failedAt: new Date(),
197 1 : failedWithError: loggable,
198 1 : }
199 : // Register as failed
200 1 : dispatchComputation({
201 1 : type: UPDATE_SCHEDULER_TASK,
202 1 : value: [task, progress]
203 1 : });
204 : // Reject for users
205 1 : task.result.reject(e);
206 1 : dispatchComputation({
207 1 : type: UNREGISTER_SCHEDULER_TASK,
208 1 : value: task,
209 1 : });
210 1 : return;
211 1 : }
212 :
213 : // Remove the task from the scheduler
214 1 : dispatchComputation({
215 1 : type: UNREGISTER_SCHEDULER_TASK,
216 1 : value: task
217 1 : });
218 1 : }
219 :
220 1 : function getTaskVariantName(task: TaskVariant): string {
221 1 : switch (task.type) {
222 0 : case COLUMN_AGGREGATION_TASK:
223 0 : return "column_aggregation";
224 1 : case TABLE_FILTERING_TASK:
225 1 : return "table_filtering";
226 0 : case TABLE_ORDERING_TASK:
227 0 : return "table_ordering";
228 0 : case TABLE_AGGREGATION_TASK:
229 0 : return "table_aggregation";
230 0 : case SYSTEM_COLUMN_COMPUTATION_TASK:
231 0 : return "system_column";
232 0 : case FILTERED_COLUMN_AGGREGATION_TASK:
233 0 : return "filtered_column_aggregation"
234 0 : case DATA_SEARCH_TASK:
235 0 : return "data_search";
236 1 : }
237 1 : }
|
