C ode G en W rangler: Data Wrangling task automation using Code-Generating Models (2025)

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.

## 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 @xcite @xcite cannot often incorporate semantics or external context (e.g., imputing city from zip code), limiting their effectiveness in complex industrial datasets. deep learning approaches @xcite @xcite 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 @xcite @xcite @xcite @xcite 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 @xcite incurs high computational costs and adds latency, making it difficult to scale for large datasets. to address these challenges, we introduce codegenwrangler, which leverages codegenerating 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 @xcite ) 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) @xcite ) by enhancing the model's capacity to retrieve context-specific knowledge to improve accuracy and relevance. proposed codegenwrangler 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)). ## background recent studies @xcite @xcite have shown that llms are capable of functioning as code generation models, which can generate code by interpreting natural language instructions @xcite @xcite , complete partially written code @xcite , and fix buggy code @xcite @xcite 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. ## method the codegenwrangler system (shown in figure 2 ) takes as input a dataset @xmath0, . . . , 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. target column ct reference tables retrieve relevant knowledge for rag column selection generate source code using code generating llm filtered dataset (d ̃) prompt construction validated code snippets filter snippets based on executability and correctness best code-snippet prompt -task instructions -external knowledge -sample data -code example prompt -task instructions -sample data -code example memory-dependent task memory-independent task dataset filtering knowledge retrieval iterative prompt optimization ## 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 @xcite . this approach involves a row-wise application of llms. the second baseline replicated the (li and döhmen, 2024) system without external memory (memoryindependent module), as outlined in @xcite , 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 mixtral-8x7b models, in alignment with the results reported in @xcite . in contrast, the codegenwrangler 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-70b-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 @xcite . for error correction, 50% of the target column values were swapped with entries from other rows to simulate realistic errors. ## results and analysis we compared performance between codegenwrangler and baselines on various datasets across di, ed and ec tasks. in effectively incorporating external data improves performance on knowledge-dependent tasks: utilizing external memory to improve performance and consistency by providing a reliable if val _accuracies > 0.9 then 19: 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 flipkart_ecommerce 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. ## 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 rowwise 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.