# Contextualized Data-Wrangling Code Generation in Computational Notebooks

arXiv: https://arxiv.org/abs/2409.13551 · HTML: https://mlsanity.com/paper/2409.13551/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 58 of its 100 references that are in the map.
- First posted on arXiv 20 Sep 2024.

## Builds on (most cited first, 30 of 58)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Program Synthesis with Large Language Models](https://mlsanity.com/paper/2108.07732/) (2021)
- [The Pile: An 800GB Dataset of Diverse Text for Language Modeling](https://mlsanity.com/paper/2101.00027/) (2020)
- [Code Llama: Open Foundation Models for Code](https://mlsanity.com/paper/2308.12950/) (2023)
- [Is Your Code Generated by ChatGPT Really Correct? Rigorous Evaluation of Large Language Models for Code Generation](https://mlsanity.com/paper/2305.01210/) (2023)
- [CodeBERT: A Pre-Trained Model for Programming and Natural Languages](https://mlsanity.com/paper/2002.08155/) (2020)
- [DeepSeek-Coder: When the Large Language Model Meets Programming -- The Rise of Code Intelligence](https://mlsanity.com/paper/2401.14196/) (2024)
- [Spider: A Large-Scale Human-Labeled Dataset for Complex and Cross-Domain Semantic Parsing and Text-to-SQL Task](https://mlsanity.com/paper/1809.08887/) (2018)
- [CodeT5: Identifier-aware Unified Pre-trained Encoder-Decoder Models for Code Understanding and Generation](https://mlsanity.com/paper/2109.00859/) (2021)
- [Measuring Coding Challenge Competence With APPS](https://mlsanity.com/paper/2105.09938/) (2021)
- [CodeGen: An Open Large Language Model for Code with Multi-Turn Program Synthesis](https://mlsanity.com/paper/2203.13474/) (2022)
- [CodeXGLUE: A Machine Learning Benchmark Dataset for Code Understanding and Generation](https://mlsanity.com/paper/2102.04664/) (2021)
- [Seq2SQL: Generating Structured Queries from Natural Language using Reinforcement Learning](https://mlsanity.com/paper/1709.00103/) (2017)
- [Unified Language Model Pre-training for Natural Language Understanding and Generation](https://mlsanity.com/paper/1905.03197/) (2019)
- [GraphCodeBERT: Pre-training Code Representations with Data Flow](https://mlsanity.com/paper/2009.08366/) (2020)
- [CodeBLEU: a Method for Automatic Evaluation of Code Synthesis](https://mlsanity.com/paper/2009.10297/) (2020)
- [Unified Pre-training for Program Understanding and Generation](https://mlsanity.com/paper/2103.06333/) (2021)
- [UniXcoder: Unified Cross-Modal Pre-training for Code Representation](https://mlsanity.com/paper/2203.03850/) (2022)
- [The Stack: 3 TB of permissively licensed source code](https://mlsanity.com/paper/2211.15533/) (2022)
- [Efficient Training of Language Models to Fill in the Middle](https://mlsanity.com/paper/2207.14255/) (2022)
- [Learning to Mine Aligned Code and Natural Language Pairs from Stack Overflow](https://mlsanity.com/paper/1805.08949/) (2018)
- [Natural Language to Code Translation with Execution](https://mlsanity.com/paper/2204.11454/) (2022)
- [Mapping Language to Code in Programmatic Context](https://mlsanity.com/paper/1808.09588/) (2018)
- [Jigsaw: Large Language Models meet Program Synthesis](https://mlsanity.com/paper/2112.02969/) (2021)
- [ReCode: Robustness Evaluation of Code Generation Models](https://mlsanity.com/paper/2212.10264/) (2022)
- [NL2Bash: A Corpus and Semantic Parser for Natural Language Interface to the Linux Operating System](https://mlsanity.com/paper/1802.08979/) (2018)

## Built on by (most cited first, 3 of 3)

- [LLM/Agent-as-Data-Analyst: A Survey](https://mlsanity.com/paper/2509.23988/) (2025)
- [Data Analysis in the Era of Generative AI](https://mlsanity.com/paper/2409.18475/) (2024)
- [CRABS: A syntactic-semantic pincer strategy for bounding LLM interpretation of Python notebooks](https://mlsanity.com/paper/2507.11742/) (2025)

## Abstract (arXiv)

Data wrangling, the process of preparing raw data for further analysis in computational notebooks, is a crucial yet time-consuming step in data science. Code generation has the potential to automate the data wrangling process to reduce analysts' overhead by translating user intents into executable code. Precisely generating data wrangling code necessitates a comprehensive consideration of the rich context present in notebooks, including textual context, code context and data context. However, notebooks often interleave multiple non-linear analysis tasks into linear sequence of code blocks, where the contextual dependencies are not clearly reflected. Directly training models with source code blocks fails to fully exploit the contexts for accurate wrangling code generation. To bridge the gap, we aim to construct a high quality datasets with clear and rich contexts to help training models for data wrangling code generation tasks. In this work, we first propose an automated approach, CoCoMine to mine data-wrangling code generation examples with clear multi-modal contextual dependency. It first adopts data flow analysis to identify the code blocks containing data wrangling codes. Then, CoCoMine extracts the contextualized datawrangling code examples through tracing and replaying notebooks. With CoCoMine, we construct CoCoNote, a dataset containing 58,221 examples for Contextualized Data-wrangling Code generation in Notebooks. To demonstrate the effectiveness of our dataset, we finetune a range of pretrained code models and prompt various large language models on our task. Furthermore, we also propose DataCoder, which encodes data context and code&textual contexts separately to enhance code generation. Experiment results demonstrate the significance of incorporating data context in data-wrangling code generation and the effectiveness of our model. We release code and data at url...