diff --git a/src/vs/workbench/services/languageRuntime/common/languageRuntimeDataExplorerClient.ts b/src/vs/workbench/services/languageRuntime/common/languageRuntimeDataExplorerClient.ts index 1592782a4f7..bc26fc96d87 100644 --- a/src/vs/workbench/services/languageRuntime/common/languageRuntimeDataExplorerClient.ts +++ b/src/vs/workbench/services/languageRuntime/common/languageRuntimeDataExplorerClient.ts @@ -402,9 +402,9 @@ export class DataExplorerClientInstance extends Disposable { this._asyncTasks.set(callbackId, promise); await this._backendClient.getColumnProfiles(callbackId, profiles, this._profileFormatOptions); - const timeout = 10000; + const timeout = 60000; - // Don't leave unfulfilled promise indefinitely; reject after 10 seconds pass + // Don't leave unfulfilled promise indefinitely; reject after one minute // for now just in case setTimeout(() => { // If the promise has already been resolved, do nothing. diff --git a/src/vs/workbench/services/positronDataExplorer/common/tableSummaryCache.ts b/src/vs/workbench/services/positronDataExplorer/common/tableSummaryCache.ts index d27e32efa79..d98b8f57b27 100644 --- a/src/vs/workbench/services/positronDataExplorer/common/tableSummaryCache.ts +++ b/src/vs/workbench/services/positronDataExplorer/common/tableSummaryCache.ts @@ -319,109 +319,134 @@ export class TableSummaryCache extends Disposable { const histogramSupported = this.isHistogramSupported(); const frequencyTableSupported = this.isFrequencyTableSupported(); - // Load the column profiles. - const columnProfileResults = await this._dataExplorerClientInstance.getColumnProfiles( - columnIndices.map((column_index): ColumnProfileRequest => { - // Get the column schema. - const columnSchema = this._columnSchemaCache.get(column_index); - - // Build the array of column profiles to load. Always load the null count. - const profiles: ColumnProfileSpec[] = [{ - profile_type: ColumnProfileType.NullCount - }]; - - // Determine whether the column is expanded. - const columnExpanded = this._expandedColumns.has(column_index); - - // If the column is expanded, load the summary stats. - if (columnExpanded) { - profiles.push({ profile_type: ColumnProfileType.SummaryStats }); - } + const columnRequests = columnIndices.map((column_index): ColumnProfileRequest => { + // Get the column schema. + const columnSchema = this._columnSchemaCache.get(column_index); + + // Build the array of column profiles to load. Always load the null count. + const profiles: ColumnProfileSpec[] = [{ + profile_type: ColumnProfileType.NullCount + }]; + + // Determine whether the column is expanded. + const columnExpanded = this._expandedColumns.has(column_index); + + // If the column is expanded, load the summary stats. + if (columnExpanded) { + profiles.push({ profile_type: ColumnProfileType.SummaryStats }); + } - // Determine whether to load the histogram or the frequency table for the column. - switch (columnSchema?.type_display) { - // Number. - case ColumnDisplayType.Number: { - // If histograms are supported, load them. - if (histogramSupported) { - // Load the small histogram. + // Determine whether to load the histogram or the frequency table for the column. + switch (columnSchema?.type_display) { + // Number. + case ColumnDisplayType.Number: { + // If histograms are supported, load them. + if (histogramSupported) { + // Load the small histogram. + profiles.push({ + profile_type: ColumnProfileType.SmallHistogram, + params: { + method: ColumnHistogramParamsMethod.FreedmanDiaconis, + num_bins: SMALL_HISTOGRAM_NUM_BINS, + } + }); + + // If the column is expanded, load the large histogram. + if (columnExpanded) { profiles.push({ - profile_type: ColumnProfileType.SmallHistogram, + profile_type: ColumnProfileType.LargeHistogram, params: { method: ColumnHistogramParamsMethod.FreedmanDiaconis, - num_bins: SMALL_HISTOGRAM_NUM_BINS, + num_bins: LARGE_HISTOGRAM_NUM_BINS, } }); - - // If the column is expanded, load the large histogram. - if (columnExpanded) { - profiles.push({ - profile_type: ColumnProfileType.LargeHistogram, - params: { - method: ColumnHistogramParamsMethod.FreedmanDiaconis, - num_bins: LARGE_HISTOGRAM_NUM_BINS, - } - }); - } } - break; } + break; + } - // Boolean. - case ColumnDisplayType.Boolean: { - // If frequency tables are supported, load them. - if (frequencyTableSupported) { - // Load the small frequency table. Note that we do not load the large - // frequency table because there are only two possible values. - profiles.push({ - profile_type: ColumnProfileType.SmallFrequencyTable, - params: { - limit: 2 - } - }); + // Boolean. + case ColumnDisplayType.Boolean: { + // If frequency tables are supported, load them. + if (frequencyTableSupported) { + // Load the small frequency table. Note that we do not load the large + // frequency table because there are only two possible values. + profiles.push({ + profile_type: ColumnProfileType.SmallFrequencyTable, + params: { + limit: 2 + } + }); - } - break; } + break; + } - // String. - case ColumnDisplayType.String: { - // If frequency tables are supported, load them. - if (frequencyTableSupported) { - // Load the small frequency table. + // String. + case ColumnDisplayType.String: { + // If frequency tables are supported, load them. + if (frequencyTableSupported) { + // Load the small frequency table. + profiles.push({ + profile_type: ColumnProfileType.SmallFrequencyTable, + params: { + limit: SMALL_FREQUENCY_TABLE_LIMIT + } + }); + + // If the column is expanded, load the large frequency table. + if (columnExpanded) { profiles.push({ - profile_type: ColumnProfileType.SmallFrequencyTable, + profile_type: ColumnProfileType.LargeFrequencyTable, params: { - limit: SMALL_FREQUENCY_TABLE_LIMIT + limit: LARGE_FREQUENCY_TABLE_LIMIT } }); - - // If the column is expanded, load the large frequency table. - if (columnExpanded) { - profiles.push({ - profile_type: ColumnProfileType.LargeFrequencyTable, - params: { - limit: LARGE_FREQUENCY_TABLE_LIMIT - } - }); - } } - break; } + break; } + } - // Return the column profile request. - return { column_index, profiles }; - }) - ); + // Return the column profile request. + return { column_index, profiles }; + }); - // Cache the column profiles that were returned. - for (let i = 0; i < columnProfileResults.length; i++) { - this._columnProfileCache.set(columnIndices[i], columnProfileResults[i]); - } + const tableState = await this._dataExplorerClientInstance.getBackendState(); - // Fire the onDidUpdate event. - this._onDidUpdateEmitter.fire(); + // For more than 10 million rows, we request profiles one by one rather than as a batch for + // better responsiveness + const BATCHING_THRESHOLD = 5_000_000; + if (tableState.table_shape.num_rows > BATCHING_THRESHOLD) { + const BATCH_SIZE = 4; + for (let i = 0; i < columnIndices.length; i += BATCH_SIZE) { + // Get the next batch of up to 4 requests + const batchColumnRequests = columnRequests.slice(i, i + BATCH_SIZE); + const batchColumnIndices = columnIndices.slice(i, i + BATCH_SIZE); + + // Send the batch of requests to getColumnProfiles + const results = await this._dataExplorerClientInstance.getColumnProfiles(batchColumnRequests); + + // Cache the returned column profiles for each index in the batch + for (let j = 0; j < results.length; j++) { + this._columnProfileCache.set(batchColumnIndices[j], results[j]); + } + + // Fire the onDidUpdate event so things update as soon as they are returned + this._onDidUpdateEmitter.fire(); + } + } else { + // Load the column profiles as a batch + const columnProfileResults = await this._dataExplorerClientInstance.getColumnProfiles( + columnRequests + ); + // Cache the column profiles that were returned. + for (let i = 0; i < columnProfileResults.length; i++) { + this._columnProfileCache.set(columnIndices[i], columnProfileResults[i]); + } + // Fire the onDidUpdate event. + this._onDidUpdateEmitter.fire(); + } } /**