Computational notebooks empower data scientists to explore data, perform analytics, and share their findings. During data exploration, the scientist uses notebooks to construct and refine data pipelines that process data in multiple stages. Extracting pipelines from a given notebook is useful in understanding the notebook’s semantics and in migrating it to production systems. However, the nature of the data exploration process, and the lack of sufficient documentation in the notebook, present two challenges in extracting the pipelines. First, notebook cells can be executed in any order, making it difficult to capture the data flow between pipeline stages. Second, data transformation operations belonging to a stage may not be cleanly separated, making it difficult to extract cohesive pipeline components.

In this paper, we propose NB2P, a novel system that automatically extracts data science pipelines from notebooks. Given an input notebook, NB2P first parses it into an Abstract Syntax Tree (AST). It then performs analysis on the AST to recover the execution order, thereby addressing the first challenge above. Next, it groups the data operations into pipeline stages based on their semantics. This step is called semantic segmentation, and it addresses the second challenge using a tree-based, learned encoding-decoding algorithm that captures the data flow and fine-grained hierarchical information in the notebook. Finally, NB2P assembles the stages and constructs the final pipeline that can be deployed into production systems.

We train NB2P on a large notebook corpus from Kaggle. We compare NB2P against baselines that use state-of-the-art large language models and other machine learning models for source code segmentation. The experimental results show that NB2P consistently outperforms the baselines while incurring low overhead.