NBSearch: Semantic Search and Visual Exploration of Computational Notebooks Xingjun Li Yuanxin Wang Hong Wang University of Waterloo University of Waterloo Uber Technologies Waterloo, ON Waterloo, ON San Francisco, CA [email protected] [email protected] [email protected] Yang Wang Jian Zhao Uber Technologies University of Waterloo San Francisco, CA Waterloo, ON [email protected] [email protected] ABSTRACT examples, support, and documentation, such as which code APIs to Code search is an important and frequent activity for developers use and how to use them, is one vital activity for developers using using computational notebooks (e.g., Jupyter). The flexibility of notebooks, as well as other development tools [2, 13, 38, 39]. A notebooks brings challenges for effective code search, where classic notebook contains a collection of cells, including markdown cells for search interfaces for traditional software code may be limited. In documentation and code cells for scripts. Unlike traditional software this paper, we propose, NBSearch, a novel system that supports code, these cells can be written and executed in any order, offering semantic code search in notebook collections and interactive vi- great convenience for developers. sual exploration of search results. NBSearch leverages advanced However, these unique features of notebooks increase the com- machine learning models to enable natural language search queries plexity of code and its management [13, 48], thus imposing chal- and intuitive visualizations to present complicated intra- and inter- lenges for code search. Developers are often unknown about which notebook relationships in the returned results. We developed NB- APIs to use [43], e.g., “how can I plot a bar chart in Python?” Seman- Search through an iterative participatory design process with two tic code search (e.g., based on module descriptions and data flows experts from a large software company. We evaluated the models among APIs used by developers [32]) could address this issue by with a series of experiments and the whole system with a controlled enabling natural language queries [50]; however, existing methods user study. The results indicate the feasibility of our analytical may not work appropriately because the structure of a notebook is pipeline and the effectiveness of NBSearch to support code search nonlinear (with loosely-connected cells), unlike a traditional code in large notebook collections. module. Cells in a notebook can be relatively independent and a group of cells can serve as a functional unit; meanwhile, they are CCS CONCEPTS interrelated, sharing data, variables, and function calls. We could treat each code cell as a basic search unit; however, there may • Human-centered computing ! Information visualization; • not be enough semantic information (e.g., comments, docstrings, Information systems ! Search interfaces; • Applied comput- etc.) available to leverage. Further, to make sense of the search re- ing ! Document searching. sults, advanced user interfaces, rather than simple ranked lists, are KEYWORDS needed [51], because of the complex relationships between returned cells and notebooks. Semantic code search, search result visualization, computational To address these challenges, we propose a novel search interface, notebooks, document and text analysis. NBSearch, that provides effective code search in an extensive collec- ACM Reference Format: tion of notebooks and visual exploration of search results. As a first arXiv:2102.01275v1 [cs.HC] 2 Feb 2021 Xingjun Li, Yuanxin Wang, Hong Wang, Yang Wang, and Jian Zhao. 2021. step, we design and develop NBSearch as a standalone tool, rather NBSearch: Semantic Search and Visual Exploration of Computational Note- than integrating into the Jupyter Notebook environment. We adopt books. In CHI ’21: ACM CHI Conference on Human Factors in Computing the semantic search approach that has been demonstrated better Systems, May 8–13, 2021, Yokohama, Japan. ACM, New York, NY, USA, performance and user experience [32]. By leveraging deep learning 14 pages. https://doi.org/10.1145/3411764.3445048 techniques, we enable this approach with individual notebook cells 1 INTRODUCTION as basic search units, to accommodate the flexible nature of note- books. Specifically, with a large notebook repository [6], we train a Computational notebooks (e.g., Jupyter) combine code, documenta- translation model to identify mappings between existing code and tion, and outputs (e.g., tables, charts, and images) within a single text descriptions (e.g., comments) in code cells. We then use the document, which provides expressive support for developers, such model to generate text descriptors for code without any comments, as adjusting code parameters to see results on the fly. Searching so each cell is associated with code and descriptors, which became CHI ’21, May 8–13, 2021, Yokohama, Japan our search database. Next, we build a language model that charac- 2021. ACM ISBN 978-1-4503-8096-6/21/05...$15.00 terizes all the cell descriptors in a learned semantic space. When https://doi.org/10.1145/3411764.3445048 a user initiates a query in natural language, NBSearch utilizes the CHI ’21, May 8–13, 2021, Yokohama, Japan Xingjun Li, Yuanxin Wang, Hong Wang, Yang Wang, and Jian Zhao language model to process the query and retrieve search results (i.e., development. For example, messes in code may accrue during the a collection of relevant cells from multiple notebooks) by searching process, and developers may lose track of their thought processes. in the semantic space of descriptors based on similarity and then To help track code history, several tools have been proposed, such as retrieving associated code segments based on the code-descriptor Variolite [21] and Verdant [22], which support fast versioning and pairs. To facilitate the exploration of search results, we design a visual exploration of history. Code Gathering Tools [13] help find, user interface that exhibits many aspects (e.g., variables used in clean, recover, and compare versions of code in Jupyter Notebooks. code) in the search results. Specifically, it features an interactive vi- Further, BURRITO [12] and TRACTUS [46] provide the provenance sualization, NBLines, which reveals both intra- and inter-notebook of analysis workflows by capturing and displaying code outputs, relationships among cells. development activities, users’ hypotheses, etc. More recently, the NBSearch is developed through a user-centered, iterative, par- flexibility of notebook-style programming has been explored for ticipatory design process with two senior developers at Uber (who providing code tutorials [14]. are also the authors). Over the course of six months, we main- Nevertheless, the above tools for notebooks assume developers tained regular meetings and consultations with the experts on mul- produce the code, neglecting the support for code search activities tiple communication platforms (e.g., video conferencing, Slack, and that occur quite frequently in software development [39]. Develop- emails), in which several prototypes were created and tested (e.g., ers have to use generic search engines (e.g., Google Search) or those sketches, low-fidelity, and high-fidelity prototypes). We evaluated built for traditional code repositories (e.g., SourceForge), whose the machine learning models of NBSearch by running comprehen- effectiveness is unknown due to the uniqueness of notebooks. In sive experiments to compare different model alternatives, using this paper, we focus on supporting code search for computational the public notebook repository [6]. The results indicate that our notebooks to improve developers’ efficiency in identifying desired analytical pipeline is promising and reveal deep insights into se- code samples and APIs to use. lecting models for building the pipeline. Moreover, we assessed the system as a whole by conducting user studies with 12 professional 2.2 Code Search and Visualization software engineers and data scientists at Uber by using their inter- There exist several code search tools, in addition to generic search nal notebook collections. The results show that NBSearch is useful engines such as Google. Sim et al. [43] comprehensively compared and effective for supporting participants’ code search tasks when Koders [49], Google, Google Code Search, Krugle [25], and Source- programming with notebooks. The code associated with this paper Forge [44], with different sizes of search target and motivation for can be accessed at https://github.com/WatVis/NBSearch. search. These search engines index code based on API keywords In summary, our contributions in this paper include: and method signatures, which may have limitations because code • An analytical pipeline that enables the semantic search of cells in notebooks are organized quite freely and the signatures are code in computational notebook collections with the power difficult to extract. Other approaches have been explored to pro- of deep learning techniques; vide more semantic features, allowing developers to search more • An interactive system, NBSearch, that allows developers to effectively. Examplar [32] ranks search results using application de- conduct code search and browse search results enhanced by scriptions, APIs used, and data-flows in code. CodeGenie [28] lever- a visualization widget, NBLines; and ages a test-driven approach to search reusable code. Sourcerer [29] • Experiments for both the system’s analytical pipeline and uses fine-grained structural