Venue: NAACL 2025

# CodeGenWrangler: Data Wrangling task automation using

Code-Generating Models

Ashlesha Akella

IBM Research, India ashlesha.akella@ibm.com &Abhijit Manatkar

IBM Research, India abhjiitmanatkar@ibm.com Krishnasuri Narayanam

IBM Research, India knaraya3@in.ibm.com &Sameep Mehta

IBM Research, India sameephetha@in.ibm.com

###### Abstract

Assuring the data quality of tabular datasets is essential for the efficiency of the diverse tabular downstream tasks (like summarization and fact-checking). Data-wrangling tasks effectively address the challenges associated with structured data processing to improve the quality of tabular data. Traditional statistical methods handle numeric data efficiently but often fail to understand the semantic context of the textual data in tables. Deep learning approaches are resource-intensive, requiring task and dataset-specific training. Addressing these shortcomings, we present an automated system that leverages LLMs to generate executable code for data-wrangling tasks like missing value imputation, error detection, and error correction. Our system aims to identify inherent patterns in the data while leveraging external knowledge, effectively addressing both memory-independent and memory-dependent tasks.

## 1 Introduction

Tabular datasets in industrial settings frequently encompass extensive data with numerous rows and columns. Given the pivotal role of this data in informed business decision-making (via exercising diverse tabular downstream tasks), maintaining high data quality has become increasingly crucial. Data wrangling tasks (like imputing missing values or correcting errors) are vital in enhancing the quality of tabular datasets. Such tasks require both statistical insights and domain-specific semantic understanding. Statistical methods Van Buuren (2018); Gong et al. (2021); Thomas and Rajabi (2021) cannot often incorporate semantics or external context (e.g., imputing city from zip code), limiting their effectiveness in complex industrial datasets. Deep learning approaches Lin et al. (2022); Samad et al. (2022); Huang et al. (2024) can capture intricate patterns but require dataset-specific training, which is computationally expensive for large datasets.

Large language models (LLMs) offer new potential for data wrangling Iida et al. (2021); Narayan et al. (2022); Huh et al. (2023); Jaimovitch-Lopez et al. (2023); Liu et al. (2023b, 2024); Ashlesha et al. (2024); Li and Dohmen (2024) tasks by leveraging broad contextual knowledge. Trained on extensive datasets, these models hold vast knowledge that enables contextual insights and supports semantically informed data wrangling. However, the need to invoke LLM inference calls independently for each row Narayan et al. (2022) incurs high computational costs and adds latency, making it difficult to scale for large datasets.

To address these challenges, we introduce CodeGenWrapper, which leverages code-generating LLMs for efficient data wrangling. Tabular datasets often contain inherent patterns with dependencies between specific columns. Our system identifies such data patterns, represents them as concisely formulated rules, and translates them into executable code for data wrangling tasks to enhance scalability by eliminating the need for row-level LLM inference calls.

While existing study Li and Dohmen (2024) has demonstrated the efficacy of code-generating LLMs in translating data patterns into executable code for data-wrangling tasks, their system is constrained by the language model's outdated knowledge and lack of the ability to incorporate external or domain-specific enterprise data. This can be effectively addressed using Retrieval-Augmented Generation (RAG) Lewis et al. (2020); Liu et al. (2023a) by enhancing the model's capacity to retrieve context-specific knowledge to improve accuracy and relevance.

Proposed CodeGenWrapper system employs a tailored prompt design and two pipeline

routes--one external memory-dependent (to integrate relevant external knowledge), the other memory-independent. An iterative refinement process further optimizes the generated code, addressing challenges such as efficiently selecting sample data for prompts. Later sections describe the full technical details of our proposed system (and an overview of our system demonstration is available at Ashlesha and Narayanam (2025)).

## 2 Background

Recent studies Wang and Chen (2023); Zan et al. (2023); Jiang et al. (2024) have shown that LLMs are capable of functioning as code generation models, which can generate code by interpreting natural language instructions Jiang et al. (2022); Wang et al. (2023); Dong et al. (2024), complete partially written code Barke et al. (2023); Guo et al. (2023), and fix buggy code Fan et al. (2023); Joshi et al. (2023); Zhang et al. (2024) due to their extensive training on vast source code data. However, we sought to investigate if these models could also recognize logical patterns in the data without requiring explicit descriptions to determine their potential for handling data-wrangling tasks. These models when prompted with sample data and instructions, we observed that their generated code aligned with the inherent patterns in the sample data (Figure 1). However, leveraging code-generating LLMs to automate data wrangling presents several challenges: (i) addressing tasks that depend on external or enterprise-specific knowledge beyond the dataset for accuracy (ii) correctly handling complex patterns in the data that go beyond simple one-to-one mappings requires coherent integration of different control flows in the code (iii) providing optimal data samples in prompts to ensure comprehensive coverage of data patterns (iv) determining which columns of the given dataset should be presented to the LLM for effective performance on specific wrangling tasks. Section 3 explains how our system addresses these challenges.

## 3 Method

The CodeGenWrangler system (shown in Figure 2) takes as input a dataset \(D=[c_{1},\ldots,c_{n}]\), where each \(c_{i}\) is an attribute (column) of the dataset, a target column \(c_{T}\), and a data wrangling task, such as data imputation (DI), error detection (ED), or error correction (EC).

For DI, the task is to predict the missing values of the dataset column \(D[c_{T}]\). For ED, the task is to identify the erroneous entries in \(D[c_{T}]\), and for EC, the task is to detect erroneous entries in \(D[c_{T}]\) and impute them.

Figure 1: Illustrative examples of code snippets generated by the CodeGenWrapper system, demonstrating its ability to handle data wrangling for Memory Independent (Left) and Memory Dependent (Right) tasks. A few more code snippets are shown in Appendix B





### Datasets

We used datasets from (Narayan et al., 2022; Ashlesha et al., 2024), which were collected from various sources like Kaggle1 and OpenML2. These datasets span across multiple domains and contain numerous columns and rows. Each dataset is split (as in (Ashlesha et al., 2024)) into three sets: a _train_ set for iteratively constructing and improving the prompt for obtaining the optimal code snippets, a _validation_ set for validating the performance of intermediate code snippets, and a final _test_ set where we evaluate the performance of our system.

Footnote 1: [https://www.kaggle.com](https://www.kaggle.com)

Footnote 2: [https://openml.org](https://openml.org)

### Dataset Filtering for Relevant Columns

Given \(D,c_{T}\) and a task, the system identifies relevant columns by calculating permutation importances for each column in a learned Histogram-based Gradient Boosting Classification Tree (Guryanov, 2019) for predicting the target column. The relevant columns \(\tilde{c}=[c^{*}_{1},\ldots c^{*}_{k}]\) with the highest permutation importances are selected to form a subset of the data, denoted as \(\tilde{D}=D[\widetilde{c},c_{T}]\), which contains only the relevant and target columns. This helps reduce noise and ensures the code generation LLM focuses on essential data patterns within its limited context length.

### Knowledge Retrieval

In addition to the LLM's parametric memory, the knowledge required for the LLM to generate source code can come from multiple other sources: it may be derived directly from the dataset itself (e.g., set the 24_hour_service column in the Starbucks dataset (Alice, 2017) to 'True' if the values of both the columns opening_time and closing_time are midnight), or it may come from external or enterprise datasets (e.g., mapping cities to respective states or imputing job role based on job title).

To accommodate knowledge inclusion from varied sources, our system employs two parallel modules for generating code snippets: a **memory-independent module**, which relies solely on patterns derived directly from the dataset, and a **memory-dependent module**, which incorporates relevant contextual knowledge from the external knowledge base \(KB=\{T_{1},T_{2},\dots,T_{m}\}\), where each \(T_{i}\) is a tabular data. To retrieve the relevant knowledge, we compute semantic similarity between the sample rows of dataset \(D\) and each table \(T_{i}\in KB\).

Let the embedding of a row \(r\) of any dataset be denoted by \(\mathbf{e}_{r}\), computed as:

\[\mathbf{e}_{r}=\text{concat}(\mathbf{h}^{1}_{r},\mathbf{h}^{2}_{r},\dots, \mathbf{h}^{n}_{r}),\]

where \(\mathbf{h}^{j}_{r}=\text{LLM}(r[c_{j}])\) is the hidden state computed by the encoder-only language model (all-miniLM-L6-v2 in our setup) for the \(j^{th}\) col

Figure 2: Dataset Filtering and Knowledge Retrieval of CodeGenWrangler system extract relevant information before it automatically generates a few code snippets in iterations that collectively capture the data wrangling task.



umn \(c_{j}\) of the row \(r\). The similarity score \(sim(T_{i})\) for a table \(T_{i}\) is computed as:

\[sim(T_{i})=\sum_{\begin{subarray}{c}r_{D}\in D\\ r_{T_{i}}\in T_{i}\end{subarray}}\mathbf{e}^{\intercal}_{r_{D}}\mathbf{e}_{r_{ T^{l}}},\]

where \(\mathbf{e}_{r_{D}}\) and \(\mathbf{e}_{r_{T_{i}}}\) are the embeddings of the rows \(r_{D}\) (sampled from \(D\)) and \(r_{T_{i}}\) (sampled from \(T_{i}\)) respectively. We select the top-\(k\) tables \(\mathcal{T}=\{T^{*}_{1},\ldots,T^{*}_{k}\}\) such that the similarity score \(sim(T^{*}_{i})\) exceeds a fixed threshold.

Further, the **memory-independent module** has two types of tasks: **row-level tasks**, which use only the data in the current row to generate code (e.g., imputing the 24_hour_service column using opening_time and closing_time columns), and **exemplar-based tasks**, where patterns are inferred from a small set of examples in the prompt.

### Prompt Construction

For each module above, code is generated by prompting a code-generating LLM, requiring a narrowly tailored prompt structure. The prompt consists of the following components. **Task instructions**: contains a description of the task to instruct the LLM to detect patterns in the data and write a Python function corresponding to the task. **External knowledge** (reference tables): with _memory-dependent tasks_, a set of rows retrieved from relevant tables from the external knowledge base. **Sample data**: a small subset of rows sampled from the dataset. For _exemplar-based tasks_, a few additional rows from the dataset similar to each sampled row are also added alongside each of the sampled rows, enabling the LLM to infer context and patterns effectively. **Code example**: the latest and most effective code snippet generated. Figure 3 provides an example of the prompt structure.

### Sample data for the Prompt

The system employs an unsupervised clustering approach to select diverse rows of the dataset for inclusion in the prompt. Given a training dataset split \(\tilde{D}_{train}\) (containing only the relevant columns), the process involves the following steps.

For each row \(r\in\tilde{D}_{train}\), an embedding \(\mathbf{e}_{r}\) is computed as described in Section 3.3. The set of embeddings \(\{\mathbf{e}_{r}\}\) is partitioned into \(k\) clusters using k-means clustering. Each cluster is represented by its centroid \(\mathbf{c}_{i}\) (\(i\in\{1,2,\dots,k\}\)). For each cluster \(\mathcal{C}_{i}\), the row embedding \(\mathbf{e}_{r^{*}}\) closest to the centroid \(\mathbf{c}_{i}\) is selected as the representative sample:

\[r^{*}=\arg\min_{\mathbf{e}_{r}\in\mathcal{C}_{i}}\|\mathbf{e}_{r}-\mathbf{c}_{ i}\|^{2}.\]

The corresponding row \(r^{*}\) is then included in the prompt as the sample data. For exemplar-based tasks, along with each row \(r^{*}\), a set of rows similar to \(r^{*}\) (based on semantic similarity of embeddings) from \(\tilde{D}_{train}\) is included as additional examples in the prompt. The resulting set of representative rows and additional examples ensures semantic diversity and relevance for the sample data while effectively covering the training data.

### Iterative Prompt Optimization

The system employs an iterative approach Wang et al. (2022) for both memory-dependent and memory-independent modules. Algorithm 1 outlines the process of optimizing prompts iteratively, incorporating the best-performing code snippets from previous iterations. At each iteration, prompts are built using different chunks of the train set, and multiple outputs are sampled, filtering out non-executable and low-accuracy code.

This approach tackles two challenges: first, the iterative process incorporates diverse data samples, generating a set of code snippets that collectively capture various patterns. This eases the need for a single perfect snippet. Second, by including the best-performing snippet from previous iterations, the prompt is incrementally refined, improving code quality and task alignment.

Figure 3: Prompt template for code generation





### Utilizing multiple code snippets

The system generates multiple code snippets, each independently applied to the dataset. The outputs from these snippets are evaluated for each row, and a majority voting approach is employed to determine the final output value for that row. This approach enhances our solution's robustness by bringing consensus among generated code snippets, thereby mitigating the risk of individual code snippets producing erroneous outputs and improving the overall reliability and accuracy of the data-wrangling process.

## 4 Experiments

We evaluated the CodeGenWrangler system through controlled experiments, comparing it to two baselines. The first baseline used a row-wise LLM approach for missing value imputation, error detection, and correction, following the method described in Ashlesha et al. (2024). This approach involves a row-wise application of LLMs. The second baseline replicated the Li and Dohmen (2024) system without external memory (_memory-independent module_), as outlined in Li and Dohmen (2024), which operates without leveraging an external knowledge base, distinguishing it from our proposed system.

To ensure a rigorous comparison, we employed three distinct LLM models across the experimental setups. The _row-wise_ LLM baseline leveraged results derived from flan-t5-xxl and mixtal-8x7b models, in alignment with the results reported in Ashlesha et al. (2024). In contrast, the CodeGenWrapper system, both with and without external memory, utilized state-of-the-art code models, codellama-34b-instruct Roziere et al. (2023) and deepseek-coder-33b-instruct Guo et al. (2024), selected for their relevance in handling code generation tasks. Crucially, to guarantee the validity and fairness of the evaluation, all setups incorporated llama-3.1-7bb-instruct as a common baseline model, controlling for architectural and computational differences across the experimental conditions.

For imputation and error detection, we used datasets from Ashlesha et al. (2024). For error correction, 50% of the target column values were swapped with entries from other rows to simulate realistic errors.

## 5 Results and Analysis

We compared performance between CodeGenWrangler and baselines on various datasets across DI, ED and EC tasks. In Table 1, we report results on datasets which reveal some key insights. Complete results for all datasets can be found in Appendix A.1. Broadly, we make the following observations:

**Effectively incorporating external data improves performance on knowledge-dependent tasks**: Utilizing external memory to improve performance and consistency by providing a reliable

Figure 4: Example of code generated for complex data pattern, for imputing product_category_tree

and up-to-date knowledge base, which is particularly evident for tasks like DI and EC in the Airline dataset. CodeGenWrangler efficiently uses a reference table of country and continent codes, outperforming row-level baselines and variants relying solely on LLMs' internal knowledge, which are prone to errors from hallucinations.

**Code generation is an effective strategy when the data pattern can be expressed in exact logical terms**: The generated code outperforms the row-level baseline by applying precise logic, such as comparing opening_time and closing_time for 24_hours_service in the Starbucks dataset or using the Change_in_rank sign to impute Gained_in_rank and Dropped_in_rank in the fortune1000_2023 dataset. In these cases, the row-level baseline underperforms as it lacks the ability to apply precise logical decision-making and must rely on the LLM's ability to generalize from a limited number of in-context examples.

**Generating code based on diverse data samples effectively captures complex patterns**: For the row-level baseline, models rely on a small set of in-context examples, which may not be sufficient when the data pattern is complex (e.g., determining the product category taxonomy from the name alone in the product_category_tree column of the fliptext_ecommerco dataset as shown in Figure 4). By generating multiple code snippets over a diverse set of samples, the code captures information across the dataset and distills it into concise heuristics that better represent the pattern. Although these heuristics may not guarantee perfect accuracy, this approach significantly outperforms the row-level baseline.

**Code generation is less effective on Error Detection tasks**: CodeGenWrangler competes well in DI and EC but struggles with ED due to the variety of errors, like syntactic anomalies or semantic mismatches. Such errors are difficult to capture using concise code snippets. It performs poorly on tasks like the Industry column in fortune1000_2023 but excels when errors can be captured via logical rules (e.g., 24_hours_service in starbucks) or verified with external knowledge (e.g., Airport Continent / Country Code in Airlines). We observe that the code generation is not very effec 

tive for datasets that need deep semantic understanding or probabilistic reasoning or those which do not follow clear logical patterns.

The number of LLM calls required by _row-level_ method is proportional to the number of dataset rows. In contrast, our system reduces the number of LLM calls by a factor of \(10\) approximately compared to _row-level_ method (see Figure 5).

Figure 6 demonstrates the effectiveness of Iterative Prompt Optimization, where refining prompts with the best-performing code improves alignment. High semantic complexity datasets like Ticket Type and Ticket Subject required up to 40 iterations, while simpler datasets like Airport Country Code converged in fewer than 10.

## 6 Conclusion and Future Work

We proposed a system to perform data wrangling on tabular datasets using code-generating LLMs. Our system generates source code by encoding rules that capture the logical patterns in the datasets. It generates multiple task-specific code snippets for each data pattern and chooses the best code snippet via a majority vote for higher reliability. The generated code snippets are executed to carry out data-wrangling tasks to replace the expensive row-wise LLM inference calls by the state-of-the-art approaches for scaling to large datasets. Our system can also handle memory-dependent tasks that require task-specific additional context provided as external domain knowledge. It adopts an iterative prompt refinement strategy to optimize the generated code for accuracy and efficiency. We plan to extend our approach for its applicability to other downstream tasks and different language models. We plan to evaluate the performance of the system on more realistic noisy and incomplete datasets.

## References

* Alice [2017] Alice. 2017. Starbucks Dataset. [https://data.world/alice-c/starbucks/workspace/file?filename=Starbucks+in+California.csv](https://data.world/alice-c/starbucks/workspace/file?filename=Starbucks+in+California.csv). Accessed on 3-Oct-2024.
* Ashlesha et al. [2024] Akella Ashlesha, Abhijit Manatkar, Brij Chavda, and Hima Patel. 2024. An automatic prompt generation system for tabular data tasks. In _Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics (NAACL): Human Language Technologies (Industry Track)_, pages 191-200.
* Ashlesha and Narayanam [2025] Akella Ashlesha and Krishnasuri Narayanam. 2025. Data Wrangling task automation using Code-Generating Language Models. _Proceedings of the AAAI Conference on Artificial Intelligence (AAAI)_.
* Barke et al. [2023] Shraddha Barke, Michael B. James, and Nadia Polikarpova. 2023. Grounded copilot: How programmers interact with code-generating models. _Proceedings of the ACM on Programming Languages (OOPSLA)_, 7(1):85-111.
* Cricsheet [2023] Crcisheet. 2023. IPL Matches. [https://www.kaggle.com/datasets/patrick1912/ipl-complete-dataset-20082020?select=matches.csv](https://www.kaggle.com/datasets/patrick1912/ipl-complete-dataset-20082020?select=matches.csv). Accessed on 3-Oct-2024.
* Dong et al. [2024] Yihong Dong, Xue Jiang, Zhi Jin, and Ge Li. 2024. Self-collaboration code generation via chatgrt. _ACM Transactions on Software Engineering and Methodology (TOSEM)_, 33(7):1-38.

Figure 5: Number of LLM calls (Llama–3.1-7bb-instruct) required for DI task across 12 datasets: D0-D3 (Airline), D4-D6 (Customer Support), and D7-D11 (Fortune 1000).

Figure 6: Gain with _Iterative Prompt Optimization_. (Legend format <Task Dataset: Target-col>)



Zhiyu Fan, Xiang Gao, Martin Mirchev, Abhik Roychoudhury, and Shin Hwei Tan. 2023 Automated repair of programs from large language models. In _Proceedings of the IEEE/ACM 45th International Conference on Software Engineering (ICSE)_, pages 1469-1481.
* Forbes [2019] Forbes. 2019. Fortune 1000. [https://www.kaggle.com/datasets/againlloty/fortune1000](https://www.kaggle.com/datasets/againlloty/fortune1000). Accessed on 3-Oct-2024.
* Gong et al. [2021] Yongshun Gong, Zhibin Li, Jian Zhang, Wei Liu, Yilong Yin, and Yu Zheng. 2021. Missing value imputation for multi-view urban statistical data via spatial correlation learning. _IEEE Transactions on Knowledge and Data Engineering (TKDE)_, 35(1):686-698.
* Guo et al. [2023] Daya Guo, Canwen Xu, Nan Duan, Jian Yin, and Julian Mcauley. 2023. LongCoder: A long-range pretrained language model for code completion. In _Proceedings of the 40th International Conference on Machine Learning (ICML)_, pages 12098-12107.
* Guo et al. [2024] Daya Guo, Qihao Zhu, Dejian Yang, Zhenda Xie, Kai Dong, Wentao Zhang, Guanting Chen, Xiao Bi, Yu Wu, YK Li, et al. 2024. Deepseek-coder: When the large language model meets programming-the rise of code intelligence. _arXiv preprint arXiv:2401.14196_.
* Guryanov [2019] Aleksei Guryanov. 2019. Histogram-based algorithm for building gradient boosting ensembles of piecewise linear decision trees. In _Proceedings of the 8th International Conference on Analysis of Images, Social Networks and Texts (AIST)_, pages 39-50.
* Huang et al. [2024] Buliao Huang, Yunhui Zhu, Muhammad Usman, and Huanhuen Chen. 2024. Semi-supervised learning with missing values imputation. _Knowledge-Based Systems (KBS)_, 284:111171.
* Huh et al. [2023] Joon Suk Huh, Changho Shin, and Elina Choi. 2023. Pool-search-demonstrate: Improving data-wrangling llms via better in-context examples. In _NeurIPS 2023 Second Table Representation Learning Workshop_.
* Iida et al. [2021] Hiroshi Iida, Dung Thai, Varun Manjunatha, and Mohit Iyyer. 2021. TABBIE: Pretrained representations of tabular data. In _Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics (NAACL): Human Language Technologies_, pages 3446-3456.
* Jaimovitch-Lopez et al. [2023] Gonzalo Jaimovitch-Lopez, Cesar Ferri, Jose Hernandez-Orallo, Fernando Martinez-Plumed, and Maria Jose Ramirez-Quintana. 2023. Can language models automate data wrangling? _Machine Learning (ML)_, 112(6):2053-2082.
* Jiang et al. [2022] Ellen Jiang, Edwin Toh, Alejandra Molina, Kristen Olson, Claire Kayacik, Aaron Donosbach, Carrie J Cai, and Michael Terry. 2022. Discovering the syntax and strategies of natural language programming with generative language models. In _Proceedings of the 2022 CHI Conference on Human Factors in Computing Systems (CHI)_, pages 1-19.
* Jiang et al. [2024] Juyong Jiang, Fan Wang, Jiasi Shen, Sungju Kim, and Sunghun Kim. 2024. A survey on large language models for code generation. _arXiv preprint arXiv:2406.00515_.
* Joshi et al. [2023] Harshti Josh, Jose Cambronero Sanchez, Sumit Gulwani, Vu Le, Gust Verbruggen, and Ivan Radicek. 2023. Repair is nearly generation: Multilingual program repair with llms. In _Proceedings of the AAAI Conference on Artificial Intelligence (AAAI)_, volume 37, pages 5131-5140.
* Lewis et al. [2020] Patrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Kuttler, Mike Lewis, Wen-tau Yih, Tim Rocktaschel, Sebastian Riedel, and Douwe Kiela. 2020. Retrieval-augmented generation for knowledge-intensive nlp tasks. In _Advances in Neural Information Processing Systems (NeurIPS)_, volume 33, pages 9459-9474.
* Li and Dhomen [2024] Xue Li and Till Dhomen. 2024. Towards efficient data wrangling with llvms using code generation. In _Proceedings of the Eighth Workshop on Data Management for End-to-End Machine Learning (DEEM)_, pages 62-66.
* Lin et al. [2022] Wei-Chao Lin, Chih-Fong Tsai, and Jia Rong Zhong. 2022. Deep learning for missing value imputation of continuous data and the effect of data discretization. _Knowledge-Based Systems (KBS)_, 239:108079.
* Liu et al. [2023a] Pengfei Liu, Weizhe Yuan, Jinlan Fu, Zhengbao Jiang, Hiroaki Hayashi, and Graham Neubig. 2023a. Pretrain, prompt, and predict: A systematic survey of prompting methods in natural language processing. _ACM Computing Surveys_, 55(9):1-35.
* Liu et al. [2023b] Shang-Ching Liu, ShengKun Wang, Tsungyao Chang, Wenqi Lin, Chung-Wei Hsiume, Yi-Chen Hsieh, Yu-Ping Cheng, Sian-Hong Luo, and Jianwei Zhang. 2023b. Jarvix: A llm no code platform for tabular data analysis and optimization. In _Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing (EMNLP): Industry Track_, pages 622-630.
* Liu et al. [2024] Yilun Liu, Shimin Tao, Xiaofeng Zhao, Ming Zhu, Wenbing Ma, Junhao Zhu, Chang Su, Yutai Hou, Miao Zhang, Min Zhang, et al. 2024. Coachlm: Automatic instruction revisions improve the data quality in llm instruction tuning. In _Proceedings of the IEEE 40th International Conference on Data Engineering (ICDE)_, pages 5184-5197.
* Mei et al. [2021] Yinan Mei, Shaoxu Song, Chenguang Fang, Haifeng Yang, Jingyun Fang, and Jiang Long. 2021. Capturing semantics for imputation with pre-trained language models. In _2021 IEEE 37th International Conference on Data Engineering (ICDE)_, pages 61-72.
* Narayan et al. [2022] Avanika Narayan, Ines Chami, Laurel Orr, Simran Arora, and Christopher Re. 2022. Can Foundation Models Wrangle Your Data? _Proceedings of the VLDB Endowment (PVLDB)_, 16(4):738-746 

Baptiste Roziere, Jonas Gehring, Fabian Gloeckle, Sten Sootla, Itai Gat, Xiaoqing Ellen Tan, Yossi Adi, Jingyu Liu, Romain Sauvestre, Tal Remez, et al. 2023. Code llama: Open foundation models for code. _arXiv preprint arXiv:2308.12950_.
* Samad et al. (2022) Manar D Samad, Sakib Abrar, and Norou Diawara. 2022. Missing value estimation using clustering and deep learning within multiple imputation framework. _Knowledge-Based Systems (KBS)_, 249:108968.
* Thomas and Rajabi (2021) Tressy Thomas and Enayat Rajabi. 2021. A systematic review of machine learning-based missing value imputation techniques. _Data Technologies and Applications (DTA)_, 55(4):558-585.
* Wan Buuren (2018) Stef Van Buuren. 2018. _Flexible imputation of missing data_. CRC press.
* Wang et al. (2022) Boshi Wang, Xiang Deng, and Huan Sun. 2022. Iteratively prompt pre-trained language models for chain of thought. In _Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing (EMNLP)_, pages 2714-2730.
* Wang and Chen (2023) Jianxun Wang and Yixiang Chen. 2023. A review on code generation with lms: Application and evaluation. In _Proceedings of the 2023 IEEE International Conference on Medical Artificial Intelligence (MedAI)_, pages 284-289. IEEE.
* Wang et al. (2023) Yue Wang, Hung Le, Akhilesh Gotmare, Nghi Bui, Junnan Li, and Steven Hoi. 2023. CodeT5+: Open code large language models for code understanding and generation. In _Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing (EMNLP)_, pages 1069-1088.
* Zan et al. (2023) Daoguang Zan, Bei Chen, Fengji Zhang, Dianjie Lu, Bingchao Wu, Bei Guan, Yongji Wang, and Jian-Guang Lou. 2023. Large language models meet nl2code: A survey. In _Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (ACL)_, pages 7443-7464.
* Zhang et al. (2024) Jialu Zhang, Jose Pablo Cambronero, Sumit Gulwani, Vu Le, Razica Piskac, Gustavo Soares, and Gust Verbruggen. 2024. Pydex: Repairing bugs in introductory python assignments using llvms. _Proceedings of the ACM on Programming Languages (OOPSLA)_, 8(1):1100-1124.

## Appendix A Appendix

### Comprehensive Evaluation

The full results in Table 2 present the performance of the proposed system on Error Detection and Missing Value Imputation tasks across 21 datasets each. These datasets, sourced from Ashlesha et al. (2024), originate from publicly available repositories such as Kaggle and OpenML, ensuring a diverse range of real-world data patterns. The results highlight the system's effectiveness in handling various data complexities, demonstrating consistent performance across multiple datasets and validating its adaptability to different data quality tasks.

## Appendix B Example Code snippets

This section presents the example code snippets, which illustrate different task-specific codes generated for different datasets. These snippets highlight the adaptability of our approach in capturing diverse data patterns effectively. For example, Figure 7 shows the code generated by our system on two different datasets, imputing the city column in the IPM Matches dataset Crischeet (2023), and imputing the Fortune 1000 dataset Forbes (2019). The code generated by the system is used for two tasks: (i) inputting the '24_hour_service' column in the 'Starbucks' dataset Alice (2017), and (ii) imputing the 'city' column in the 'Restaurant' dataset Mei et al. (2021). These are shown in Figure 8.

## Appendix C Example Prompts for Diverse Datasets

This section presents example prompts which are automatically constructed by the system for memory dependent tasks 9 and memory independent task 10. These prompts are designed to incorporate relevant data patterns, external knowledge (when applicable), and iterative refinements to enhance the quality of generated code snippets 



## Appendix A

\begin{table} \end{tabular}
\end{table}
Table 2: Results for extended datasets on Data Imputation and Error Detection tasks

Figure 7: Code Generated by our system 1. to impute ‘city’ column in ‘IPM Matches’ dataset (Cricsheet, 2023). 2. to impute ‘Gained in rank’ and ’Dropped in Rank‘ columns in ‘Fortune 1000’ dataset (Forbes, 2019)



**Starbucks dataset impute 24_hour_service**

```
deftask(inputs): regular_hours=inputs.get('regular hours') saturday_opening_times= inputs.get('$saturdayopenings') sunday_opening_times= inputs.get('$sundayopeningtimes') ifregutar_hours=='nan'or saturyday_opening_times=='nan'or saturyday_opening_times=='nan'returnFalse ifregutar_hours=='12:00AMtoc12:00AM' andsaturaday_opening_times=='12:00AM to12:00AM' andsanday_opening_times='12:00AM to12:00AM': returnTrue return False
```

**Restaurant dataset impute city**

```
deftask(data):ifdatal{'phone'}.startswith('415'): return'samfrancisco' elifdat['phone'].startswith('404'): return'atlandata' elifdata['phone'].startswith('213'): return'losangles' elifdat['phone'].startswith('212'): return'newyork' elifdata['phone'].startswith('718'): return'queens' else: return'Unknown'
```

**Restauural dataset impute city**

Figure 8: Code Generated by the system: (i) To impute ‘24_hour_service’ column in ‘Starbucks’ dataset (Alice, 2017) (ii) To impute ’city’ column in ‘Restaurant’ dataset (Mei et al., 2021)

Figure 9: Example prompt for Memory Dependent task.





#### Instructions:

Given a series of examples, your task is to identify the pattern between the input columns and their corresponding output values. Write a Python function named taskthat accurately maps the input data to the correct output based on this identified pattern.

### Key Guidelines:

1. Analyze Patterns:

Observe the logical relationships and string patterns within the input data. Focus on identifying consistent connections between input columns and their corresponding outputs.

2. Optimize the logic:

- Use regular expressions when necessary to match specific conditions or patterns efficiently.

- Erploy methods like startswithand endswitchinstead of generic comparisons for precise string matching.

3. Comprehensive Coverage:

- Ensure your code considers all possible patterns and conditions given in the test examples.

- Write as many ifstatements as required to handle each identified pattern thoroughly.

4. Relevant Features Only:

Utilize only the columns that show a consistent relationship to the output. Avoid introducing unnecessary complexity by including irrelevant columns.

5. Default Behavior:

If no recognizable pattern is found, the function should return \({}^{\text{t}}\)Unknown":

[FIGURE:S6.F0][ENDFIGUR

Figure 10: Example prompt for Memory Independent task.