PY-75205 code completion for column names doesn't work with large datasets. #PY-75205 Fixed

* In case we were unable to parse the column names from the standard result of the getFrameVars command, we make the second attempt to get the column names - an additional request specialized for dataframes data retrieval.

(cherry picked from commit ec0b807117934b804daeb44cd7d653a566418f0f)

GitOrigin-RevId: 563af434addc7e723ac6e90010286b12e53e6134
This commit is contained in:
Natalia.Murycheva
2025-02-24 17:21:33 +01:00
committed by intellij-monorepo-bot
parent 66449d422a
commit 077ee31039
3 changed files with 21 additions and 15 deletions

View File

@@ -14,7 +14,7 @@ import org.jetbrains.annotations.Nullable;
import java.util.List;
import java.util.Set;
import static com.jetbrains.python.debugger.values.DataFrameDebugValueUtilKt.getColumnData;
import static com.jetbrains.python.debugger.values.DataFrameDebugValueUtilKt.getInformationColumns;
public final class DataFrameDebugValue extends PyDebugValue {
private final ColumnNode treeColumns = new ColumnNode();
@@ -65,10 +65,13 @@ public final class DataFrameDebugValue extends PyDebugValue {
public void ok(String value) {
myLoadValuePolicy = ValuesPolicy.SYNC;
myValue = value;
DataFrameDebugValue.InformationColumns columns = getColumnData(value);
if (columns != null) {
setColumns(columns);
}
try {
DataFrameDebugValue.InformationColumns columns = getInformationColumns(value);
if (columns != null) {
setColumns(columns);
}
} catch (Exception ignored) {}
for (XValueNode node : myValueNodes) {
if (node != null && !node.isObsolete()) {
updateNodeValueAfterLoading(node, value, "", null);

View File

@@ -43,15 +43,13 @@ private val COLUMN_NAMES_REGEX = Regex("'(.*?(?<!\\\\))'|\"(.*?(?<!\\\\))\"")
* The [value] is produced by _get_df_variable_repr.
*/
@ApiStatus.Internal
fun getColumnData(value: @NlsSafe String): DataFrameDebugValue.InformationColumns? {
val columnData = getData(value) ?: run {
LOG.warn("Can't retrieve column data from ${value}")
return null
}
@Throws(CannotRetrieveColumnDataException::class)
fun getInformationColumns(value: @NlsSafe String): DataFrameDebugValue.InformationColumns? {
val columnsNames = getDataWithColumnsNames(value) ?: throw CannotRetrieveColumnDataException("Can't retrieve column data")
getMultiIndexData(columnData)?.let { return it }
getMultiIndexData(columnsNames)?.let { return it }
val columns = COLUMN_NAMES_REGEX.findAll(columnData).map { listOf(it.groups.filterNotNull().last().value) }.toList()
val columns = COLUMN_NAMES_REGEX.findAll(columnsNames).map { listOf(it.groups.filterNotNull().last().value) }.toList()
return if (columns.isNotEmpty()) {
DataFrameDebugValue.InformationColumns().apply {
@@ -64,7 +62,7 @@ fun getColumnData(value: @NlsSafe String): DataFrameDebugValue.InformationColumn
}
}
private fun getData(value: String): String? {
private fun getDataWithColumnsNames(value: String): String? {
var escaped = false
var inSingleQuotes = false
var inDoubleQuotes = false
@@ -101,3 +99,5 @@ private fun getMultiIndexData(columnData: @NlsSafe String): DataFrameDebugValue.
null
}
}
class CannotRetrieveColumnDataException(message: String): Exception(message)

View File

@@ -11,7 +11,7 @@ import com.jetbrains.python.debugger.ArrayChunkBuilder
import com.jetbrains.python.debugger.PyDebugValue
import com.jetbrains.python.debugger.PyFrameAccessor
import com.jetbrains.python.debugger.values.DataFrameDebugValue
import com.jetbrains.python.debugger.values.getColumnData
import com.jetbrains.python.debugger.values.getInformationColumns
fun parseVars(vars: List<DebugValue>, parent: PyDebugValue?, frameAccessor: PyFrameAccessor): XValueChildrenList {
val list = XValueChildrenList(vars.size)
@@ -27,7 +27,10 @@ fun parseVars(vars: List<DebugValue>, parent: PyDebugValue?, frameAccessor: PyFr
fun createPyDebugValue(value: DebugValue, frameAccessor: PyFrameAccessor): PyDebugValue {
return if (value.type == "DataFrame") {
val columns = value.value?.let { getColumnData(it) }
val columns = try {
value.value?.let { getInformationColumns(it) }
} catch (_: Exception) { null }
val dataFrameDebugValue = DataFrameDebugValue(value.name, value.type, value.qualifier, value.value ?: "",
value.isContainer, value.shape, value.isReturnedValue, value.isIPythonHidden,
value.isErrorOnEval,