NB2P: Generating Data Science Pipelines from Computational Notebooks

Computational notebooks empower data scientists to explore data, perform analytics, and share findings through data pipelines. Extracting these pipelines facilitates understanding and migrating them to production systems. However, notebook cells executed in arbitrary order complicate data flow between stages. Moreover, data operations may not be cleanly separated, hindering the extraction of cohesive pipeline components. In this paper, we propose NB2P, a novel system that automatically extracts data science pipelines from notebooks. Given a notebook, NB2P performs analysis at the syntax tree level to recover the execution order, then groups the data operations into pipeline stages based on semantics. It uses 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 and compare it against baselines that use state-of-the-art large language models and other approaches. The experimental results show that NB2P consistently outperforms the baselines while incurring low overhead.