Knowledge-Enhanced Program Repair for Data Science Code

TL;DR

DSrepair tackles the challenge of repairing data science code generated by LLMs by integrating knowledge graph–based API retrieval with AST-guided bug localization. It constructs DS-KG to capture API usage across seven major libraries and enriches repair prompts with API knowledge and AST-localized bug context. Empirical results on the DS-1000 benchmark show that DSrepair outperforms five baselines across four LLMs, achieving higher fix rates and lower token costs, and it demonstrates robustness to LLM non-determinism. The work contributes a DS-KG resource, an AST-based bug enrichment approach, and a generalizable retrieval-augmented repair framework that improves reliability and efficiency for data science code repair with potential extension to broader domains.

Abstract

This paper introduces DSrepair, a knowledge-enhanced program repair method designed to repair the buggy code generated by LLMs in the data science domain. DSrepair uses knowledge graph based RAG for API knowledge retrieval as well as bug knowledge enrichment to construct repair prompts for LLMs. Specifically, to enable knowledge graph based API retrieval, we construct DS-KG (Data Science Knowledge Graph) for widely used data science libraries. For bug knowledge enrichment, we employ an abstract syntax tree (AST) to localize errors at the AST node level. DSrepair's effectiveness is evaluated against five state-of-the-art LLM-based repair baselines using four advanced LLMs on the DS-1000 dataset. The results show that DSrepair surpasses all five baselines. Specifically, when compared to the second-best baseline, DSrepair demonstrates significant improvements, fixing 44.4%, 14.2%, 20.6%, and 32.1% more buggy code snippets for each of the four evaluated LLMs, respectively. Additionally, it achieves greater efficiency, reducing the number of tokens required per code task by 17.49%, 34.24%, 24.71%, and 17.59%, respectively.

Figures (7)

• Figure 1: Overview of DSrepair.
• Figure 2: Details on API KG construction (Step 1) and API Knowledge Retrieval (Step 2) in DSrepair. The code raises an error because of the mismatched array shape between np.flipud's required input and np.array_split's output. DSrepair extracts the API call in the error line and builds a SPARQL query to search the relevant RDF triples in the DS-KG, which is constructed from API documents and guided by the ontology. Finally, DSrepair maps the returned RDF triples to natural language, which will be used as a part of the repair prompt.
• Figure 3: Bug knowledge enrichment example. Stderr (standard error information) can only localize the bug at the line level, while our bug knowledge enrichment could enrich the error information to the AST node level.
• Figure 4: DSrepair prompt example. The prompt contains structural and rich information to guide LLMs for code generation.
• Figure 5: A code problem example from DS-1000. The incorrect solution is generated from Self-Repair, and the correct solution is generated from DSrepair. By incorporating knowledge of the invoked API, DSrepair can assist LLMs in generating solutions with correct API usage.