NotebookRAG: Retrieving Multiple Notebooks to Augment the Generation of EDA Notebooks for Crowd-Wisdom

TL;DR

NotebookRAG is proposed, a method that takes user intent, datasets, and existing notebooks as input to retrieve, enhance, and reuse relevant notebook content for automated EDA generation and confirms the superiority of the method in producing high-quality and intent-aligned EDA notebooks.

Abstract

High-quality exploratory data analysis (EDA) is essential in the data science pipeline, but remains highly dependent on analysts' expertise and effort. While recent LLM-based approaches partially reduce this burden, they struggle to generate effective analysis plans and appropriate insights and visualizations when user intent is abstract. Meanwhile, a vast collection of analysis notebooks produced across platforms and organizations contains rich analytical knowledge that can potentially guide automated EDA. Retrieval-augmented generation (RAG) provides a natural way to leverage such corpora, but general methods often treat notebooks as static documents and fail to fully exploit their potential knowledge for automating EDA. To address these limitations, we propose NotebookRAG, a method that takes user intent, datasets, and existing notebooks as input to retrieve, enhance, and reuse relevant notebook content for automated EDA generation. For retrieval, we transform code cells into context-enriched executable components, which improve retrieval quality and enable rerun with new data to generate updated visualizations and reliable insights. For generation, an agent leverages enhanced retrieval content to construct effective EDA plans, derive insights, and produce appropriate visualizations. Evidence from a user study with 24 participants confirms the superiority of our method in producing high-quality and intent-aligned EDA notebooks.

Figures (7)

• Figure 1: Overview of the Notebook Retrieval process. User intent (a) is mapped into EDA queries (b). Notebook markdown cells and code cells are preprocessed (1) into embeddings and components (with metadata), respectively. Then, the EDA queries are used to search embeddings for relevant chunks (c) and to guide component search (2) and enhancement (3), producing relevant components (c).
• Figure 2: Example of the dependency tracking process. Code Cells define data variables through generation and modification statements. Data Variables summarize each variable’s lifecycle and dependencies. Cell Dependencies show how variables propagate across cells. Component collects the required statements into a self-contained unit that can be executed independently. Some cases where components improve annotation accuracy (C1, C2, C3).
• Figure 3: Agent for EDA notebook generation with retrieval interface. The agent takes user intent, a dataset, and retrieval outputs (optional) as input. It begins by constructing an EDA plan (Planner), then incrementally dispatches sub-tasks (Progress Manager), generates statistical code and visualization (Coder), and integrates results (Summarizer) to update the plan, ultimately producing a complete EDA notebook. Retrieval outputs are passed through a retrieval interface to enhance both planning and code generation.
• Figure 4: A: Baseline EDA plan. B: RAG-enhanced EDA plan.
• Figure 5: A: a reused human-authored component showing the year-by-year distributions of features. B: a RAG-enhanced visualization highlighting anomalies and annotated events in the weekly average price series.