An Alternative to Cells for Selective Execution of Data Science Pipelines

TL;DR

This paper proposes replacing code cells in notebooks with variable-based actions driven by data-flow analysis, improving the management and execution of data science pipelines by reducing errors and clutter.

Contribution

It introduces a novel approach that uses variable context menus and data-flow analysis to replace code cells for selective execution in data science notebooks.

Findings

Reduces errors caused by incorrect cell execution order

Separates pipeline code from decision-making code

Automates dependency management for data science workflows

Abstract

Data Scientists often use notebooks to develop Data Science (DS) pipelines, particularly since they allow to selectively execute parts of the pipeline. However, notebooks for DS have many well-known flaws. We focus on the following ones in this paper: (1) Notebooks can become littered with code cells that are not part of the main DS pipeline but exist solely to make decisions (e.g. listing the columns of a tabular dataset). (2) While users are allowed to execute cells in any order, not every ordering is correct, because a cell can depend on declarations from other cells. (3) After making changes to a cell, this cell and all cells that depend on changed declarations must be rerun. (4) Changes to external values necessitate partial re-execution of the notebook. (5) Since cells are the smallest unit of execution, code that is unaffected by changes, can inadvertently be re-executed. To solve these issues, we propose to replace cells as the basis for the selective execution of DS pipelines. Instead, we suggest populating a context-menu for variables with actions fitting their type (like listing columns if the variable is a tabular dataset). These actions are executed based on a data-flow analysis to ensure dependencies between variables are respected and results are updated properly after changes. Our solution separates pipeline code from decision making code and automates dependency management, thus reducing clutter and the risk of making errors.

Full text

To solve these issues, we propose to replace cells as the basis for the selective execution of DS pipelines. Instead, we suggest populating a context-menu for variables with actions fitting their type (like listing columns if the variable is a tabular dataset). These actions are executed based on a data-flow analysis to ensure dependencies between variables are respected and results are updated properly after changes. Our solution separates pipeline code from decision making code and automates dependency management, thus reducing clutter and the risk of making errors.

††publicationid: pubid:

Accepted at the 45th Int. Conf on Software Engineering, NIER Track

ICSE-NIER’23, May 14–20, 2023, Melbourne, Australia

©2023 IEEE

I Introduction

Notebooks allow users to write code in cells that can be executed independently, in any order. Execution results, such as visualizations, are commonly shown close to the code cells that produced them. Code cells can be interspersed with text cells, to offer explanations or document decisions, which follows the paradigm of literate programming [1]. Overall, this makes notebooks well suited for the explorative development process of Data Science (DS) pipelines [2]. Here, developers often write small pieces of code, run them, and use the results to change existing code or decide what to implement next.

Various flavors of notebooks exist, with Jupyter Notebook [3] being most popular according to the 2021 Kaggle Survey on DS [4]. Jupyter Notebook requires language-specific kernels to execute code and to power IDE features like auto-completion, but is otherwise language-agnostic. Since Jupyter Notebook is a web application, it can run locally or be hosted as a service, like Google Colaboratory [5], Kaggle [6], or Amazon SageMaker [7]. Some integrated development environments (IDEs) like PyCharm [8] or Visual Studio Code [9] incorporate Jupyter Notebook. JupyterLab [10] is eventually meant to replace the default Jupyter Notebook GUI. The issues we discuss in this paper are independent from a specific notebook variant, however, since they stem from the core concept of notebooks: Cells.

Fig. 1 shows an example of the typical cell structure of a notebook for DS111Based on a popular notebook from Kaggle (https://www.kaggle.com/code/startupsci/titanic-data-science-solutions). using Python as the programming language. Text cells (blue) and code cells (white) alternate. The code cells (1) read a CSV file containing the training data, (2) view the dataset to gain a general understanding and to know which attributes it has, (3) separate feature vector and target, (4) configure a model (support vector machine for classification), and train it222The fit method in this example is supposed to return a new trained model rather than mutate the untrained model stored in svc.. For the sake of brevity, we omit result cells.

Based on Fig. 1 we can illustrate the problems we want to discuss in this paper:

Actual pipeline code (Cells 1, 3, 4) is mixed with value inspection code (Cell 2). The value inspection code exists solely to write more pipeline code (here Cell 3, which needs the name of the target column). Keeping value inspection cells in the notebook afterwards increases clutter [11, 12] and negatively impacts performance, since they might get re-executed unnecessarily. 2. 2.

Users can execute cells in any order. This manipulates the internal state of the notebook (e.g. contained variables and their assigned values), which is maintained throughout the entire session. However, not all orderings are correct [13, 11, 14, 15, 16, 17, 18]: For example, executing Cell 3 first on a blank state is erroneous, since this cell depends on train_df, which is computed in Cell 1. 3. 3.

Code changes partially invalidate the internal notebook state and results cells [13, 11, 14, 15, 16, 17, 18]. Say, after executing the entire notebook from Fig. 1, we add more data preparation steps to Cell 3, before assigning X_train and y_train. Now, the values of X_train and y_train in the internal notebook state are outdated, so Cell 3 must be rerun. However, trained_svc, which depends on these values, is also outdated, so Cell 4 must be rerun, too. In large notebooks, keeping track of all cells that must be rerun gets complicated, leading to developers frequently rerunning the entire notebook to be safe [16], wasting computing and development time. 4. 4.

Even without code changes, notebook state can become stale: For example, the dataset we read from disk in Cell 1 might change. Like in Problem 3, the notebook needs to be partially re-executed to account for such an event. 5. 5.

Since cells are the smallest unit of execution, even code that is unaffected by changes is re-executed. In the example from Problem 3, we rerun Cell 4 completely, although the value of svc is still valid. Calling the SVC constructor is fast, but for the same reason long-running data preparation or model training operations might be rerun unnecessarily, drastically slowing the feedback loop. This leads to an “expand then reduce” pattern [11], where developers first write small code cells, which can be executed independently, to iterate quickly and later combine them into bigger cells. This defeats the purpose of cells as a means to group logically related code together.

II State of the Art

A natural solution to ensure code gets executed in the correct order (Problem 2) and gets re-executed after changes (Problem 3) is to derive dependencies between cells using data-flow analysis and run a cell only after all cells it depends on: Dataflow notebooks [19] assign a unique identifier to cells, which stays the same even as the code in the cell changes. These IDs are used to describe the dependencies between cells. When a cell gets executed, the system ensures that all upstream dependencies are available already. Nodebook [20] keeps track of inputs and outputs of cells. Inputs are determined by parsing the code in the cell, while outputs are discovered by comparing the internal state of the notebook after the cell was run to the state before. [14] describes how data-flow analysis can be used to create a polished version of a notebook that only contains the code needed to produce the results that the user selected. NBSAFETY [17] uses data-flow analysis to detect unsafe interactions with the cells in a notebook and offer resolution advisories. To achieve this, NBSAFETY uses a mix of dynamic and static analysis. ReSplit [21] analyzes definition-usage chains between cells and within the same cell and then suggests an alternative mapping of code to cells to ensure that tightly coupled code resides in the same cell.

However, none of the existing approaches solve Problems 1, 4, and 5: They do not handle changes to values outside the notebook (Problem 4), use complete cells as the smallest unit of execution thus failing to avoid re-execution of parts unaffected by changes (Problem 5), and do not even try to address the tangling of pipeline and value inspection code (Problem 1).

III Research Questions

This brings us to the following research questions:

• •

RQ 1: How can we separate pipeline code and value inspection code?

• •

RQ 2: How can we ensure changes to external values trigger re-execution of code that depends on them?

• •

RQ 3: How can we avoid unnecessarily executing code?

RQ 1 aims to reduce clutter, RQ 2 is about correctness and RQ 3 addresses performance.

IV Approach Overview

To solve these issues, we propose to keep cells only as a means to connect code to related results and instructional or documenting text, but abandon them as the basis for the selective execution of DS pipelines. Instead, we suggest to

address RQ 1 by introducing different roles for code cells (Sec. V), or by statically typing variables and letting the development environment dynamically populate a context-menu for variables with actions that can be run on a variable of the respective type (Sec. VI), 2. 2.

address RQ 2 and RQ 3 by using static code analysis to derive a correct and minimal execution plan for a selected action (Sec. VII).

An execution plan is correct, if it contains all operations that must be executed and guarantees that each operation is executed only when all its input values are up-to-date, including external values (RQ 2). It is minimal (RQ 3), if it contains no operations that would only affect already up-to-date values or would compute values that are not accessed by other operations in the graph.

V Tags for Code Cells

For teaching materials [22] or documentation, keeping value inspection cells in the notebook can be helpful, to outline each step of the development process of a DS pipeline, including decision-making [11]. However, value inspection code is then scattered across the entire notebook, creating clutter, and re-executed each time the notebook is run, wasting time.

To separate pipeline code and value inspection code (RQ 1) in this case, we suggest to tag code cells by their purpose, as “pipeline” or “inspection”. Fig. 2 shows this for the start of the notebook from Fig. 1 with pipeline cells (white) and inspection cells (yellow). With a filter, a user can then narrow down the cells and show, say, just the pipeline cells. Only cells that are currently shown are executed. Text cells can be tagged and filtered in the same way.

VI Context-Sensitive Inspection of Values

Outside of educational notebooks, value inspections cells are often executed only once, to decide what to do next. For example, after running Cell 2 in Fig. 1 we know that the target attribute is called “Survived” and can use this knowledge to write Cell 3. Leaving Cell 2 in the notebook has little documentation value, since we already manifested the extracted information in Cell 3.

In contexts where inspection cells have no documentation effect, we suggest to avoid writing them in the first place, by additionally offering inspection actions in a context-menu for variables (Fig. 3). This way, notebooks can only contain pipeline code if desired, completely eliminating the tangling with inspection code (RQ 1).

To know which actions can be triggered on a variable, we need to know its type. Ideally, the type should be known statically, so we do not need to run code to determine which actions can be triggered on a variable. Moreover, the type should be inferred, so users can concentrate on writing their pipeline code as they are used to.

When a user selects an action, code that implements the action is executed in the background (see Sec. VII for details). Results of value inspection actions can be closed after inspection or kept outside the main flow of the notebook entirely in separate tabs, windows, or Sticky Cells [23] that float on top of the notebook and maintain their position even when the notebook is scrolled. Fig. 4 shows a mockup of the potential output for the “Show dataset” action from Fig. 3. The data is displayed in an interactive table that the user can filter (funnel icon) or sort (arrow icons). Additional actions can be triggered directly from this view, without the need to go back to the context menu, e.g. for generating a histogram of a column (chart icons).

VII Minimal Execution

Data-flow

As described in Sec. VI, we want to trigger context-sensitive actions on variables without executing unnecessary code (RQ 3). For example, if the user selects the “Show dataset” action on the variable X_train from Fig. 1, it is a waste of time to execute the entire Cell 3, since the value of y_train is never used. This requires a fine-grained data-flow graph that focuses on individual operations rather than cells. Fig. 5 shows the data-flow graph for the example notebook from Fig. 1. From this graph we can derive that we only need to evaluate the calls of read_csv and drop to compute X_train, leading to the simple execution plan in Fig. 6

Purity

We can derive an execution plan for the entire notebook up to the fit call, based on the data-flow graph from Fig. 5. The graph tells us that we need to run the entire notebook, since fit depends on all other operations. However, we can further optimize, if we know which operations are pure. A pure operations has no side effects and its outputs only depend on its parameters captured in the data-flow graph. Impure operations may read or modify external state (files), or global state (global variables or object attributes). Any pure operation can be executed in parallel to any other independent operation (that is, an operation on a different path of the dataflow graph). In contrast, independent impure operations must be executed according to their textual order. For this, additional edges that reflect the textual order are added between independent impure operations, yielding the complete execution plan. In our example, all operations are pure, except read_csv, which reads a file from disk and is impure because the file contents might change although the value of the read operation’s path parameter stays the same. Adding purity information (green background) and impurity information (red background) leads to the extended data-flow graph shown in Fig. 7. From it we can deduce that the calls to read_csv and SVC can be run in parallel, as well as the calls to drop and keep. If the calls to drop and keep where impure and the call to drop occurred textually before the one to keep, an additional edge would point in Fig. 7 from drop to keep.

Execution

The extended dataflow graph serves as an execution plan for computing the value of some variable x, when the user triggers an action (Sec. VI) on x:

Build the dataflow graph. 2. 2.

Extend it by purity information for operations. 3. 3.

Extend it by textual order edges between independent impure operations. 4. 4.

Eliminate operations that have no path to x. 5. 5.

Start executing (in any order or in parallel) nodes that have no incoming edges. After execution, delete the respective node and its outgoing edges. 6. 6.

Repeat from Step 5, until the graph is empty, at which point the value of x is available.

The first two steps ensure correctness. The first ensures that data dependencies are respected. The second preserves the order of impure operations. The third step removes all operations that have no effect on the value we want to compute, a first contribution towards minimality (RQ 3).

Re-Execution

Let us now assume the initial execution plan from Fig. 7 has been run, so the internal notebook state contains the current values of all variables. Let us further assume that the user has subsequently edited the call of the drop operation, to remove additional attributes from the feature vector after inspecting X_train. Now, if the user requests the system to update the state of the notebook, we want to ensure that all code that is affected by the change gets rerun, without re-executing code unnecessarily (RQ 3). As with any rerun, we also need to take into account potential changes to external state, which can affect impure operations (RQ 2).

Fig. 8 illustrates the affected part of the notebook, after editing the drop call. The changed operation is marked by a red border. Potentially stale variables are indicated by yellow question marks:

• •

X_train, the output of the edited drop call,

• •

train_df, the output of the impure operation read_csv – because the data-flow graph does not show the hidden dependencies of impure operations, we must always assume that something might have changed,

• •

y_train, the output of the call to keep that depends on the potentially stale value train_df.

All potentially stale values are indicated by a question mark in Fig. 8. The call to SVC is unaffected by the change and its output is still up-to-date. Therefore they are not included in the re-execution plan.

The red labels in Fig. 8 reflect the assumption that we must always rerun edited operations, impure operations, and pure operations with at least one potentially stale argument. Potentially stale arguments must be recomputed before they are used, which leads to the shown ordering.

A possibly more efficient approach is illustrated in Fig. 9. The yellow question marks on the calls to keep and fit indicate that these operations do not necessarily need re-execution, if recomputing their potentially stale arguments produces the same values as before. For this non-staleness check, we can store the previous value and compare it to the new one. If it didn’t change, we mark the argument as up-to-date. If all inputs of a pure operation are up-to-date, we do not need to re-run it and can mark its output as up-to-date. In the plan from Fig. 9 the operations drop and keep can be run in parallel after read_csv. If the dataset returned by read_csv is unchanged, we can skip the call to keep and can mark y_train as up-to-date.

This approach can eliminate expensive re-executions at the cost of remembering previous values and comparing them to new ones. For large datasets or models, such comparisons can be costly, too. Thus, the efficiency of the approach critically depends on the trade-off between re-execution time and comparison time. Even expensive comparisons might pay off, because they can eliminate re-execution of many downstream pure operations, including the followup checking of each of their results.

In this regard, an impure operation like read_csv is a particularly worthwhile target for optimization: (1) It usually occurs right at the start of a pipeline, so many other operations depend on it and must be repeated if read_csv yields new results. (2) The returned data is large, so an equality check is extremely expensive.

We can fortunately avoid executing file operations subject to a simple check and the storing of minimal extra information: The time of the last modification of each accessed file. If we call read_csv with the same path argument and the respective file has still the same last modified time, the operation returns the same result. We can, therefore, treat the last modified time internally as an extra, hidden argument, and the modified read_csv operation as pure. This principle can be generalized to other cases, e.g. a random number generator that has a fixed seed is no source of impurity anymore. This way, some often occurring cases of impurity can be eliminated, accounting for noticeable re-execution speedups. It allows us to treat all operations shown in Fig. 1 as pure.

VIII Technical Challenges and Solutions

The concept of notebooks builds on simple language-specific kernels. Our proposal requires additional kernel features, namely static typing for the context menu (Sec. VI), data-flow analysis to derive a correct but minimal execution plan (Sec. VII), and purity inference as a basis for exploiting parallelism. Implementing our proposal for arbitrary languages, would require significant additional work from kernel developers. For some languages, it would even be impossible to implement our extensions fully. Python, for instance, allows dynamically importing modules via importlib.import_module(name) [24]. Reading the module name from an encrypted file would prevent static analysis of the module.

Instead of building on an existing language, we developed Safe-DS [25], a domain specific language for implementing DS pipelines. Safe-DS can integrate existing Python libraries via a simple stub language in a statically type-safe way. This lets the compiler infer the type information that we need to provide context-specific actions on variables. The design of Safe-DS as a language for creating pipelines rather than implementing the algorithms used in the pipeline makes data-flow analysis easy, due to the lack of loops, recursion, and conditionals, which guarantees an acyclic data-flow graph. Safe-DS additionally offers annotations to mark operations as pure, empowering the compiler to prune the execution plan even in cases where purity cannot be inferred. Lastly, Safe-DS has a standard library that prefers pure operations.

IX Limitations

Generalizing our approach to notebooks written in arbitrary languages can by tricky. It depends on the quality of static analysis for types, data-flow and purity available for the respective language. If at least data-flow analysis works, one can alternatively annotate libraries with types and purity information manually. But for non-trivial libraries, this may quickly exceed available time. However, one can use a combined approach that performs a static analysis and then lets developers review and complement its results by annotating library elements with non-inferrable type or purity information [26]. Such combined tools limit the amount of manual annotation, thus making it possible to exploit the full potential of our approach in spite of the limits of static analysis.

X Future Plans

Prototype

We currently extend Safe-DS333https://github.com/Safe-DS/DSL by the discussed computation of execution graphs based on data-flow analysis and purity information. An experimental IDE for Safe-DS is already available as an extension for Visual Studio Code444https://marketplace.visualstudio.com/items?itemName=safe-ds.safe-ds. After evolving Safe-DS, we will extend its IDE by the proposed, type-based context menu for inspection of variables, including presentation of tabular datasets (Fig. 4), to eliminate inspection code from the pipeline.

Usability study

Once the prototype is complete, it will be compared quantitatively and qualitatively to Jupyter Notebook. We will measure effects on code comprehension (by separating value inspection and pipeline code), correctness (by enforcing execution of code in the right order and re-execution of changed code), and performance (by avoiding re-execution of unchanged code).

Purity Inference for Python

Moreover, we plan to add purity inference (pure, impure, unknown) to the API-Editor [26], use it to provide the unknown information, and let purity information be reflected in the Safe-DS stubs that it creates. We will also use its ability to automatically create wrappers that implement a uniform DS API on the basis of existing Python libraries.

XI Conclusions

In this paper, we presented an alternative to cells as a basis to selectively execute parts of a DS pipeline. We use static code analysis to derive a fine-grained data-flow graph that focuses on individual operations rather than cells (Sec. VII). From this, we derive an execution plan that is correct (contains all operations and runs them in a valid order) and minimal (does not contain unnecessary operations). Information about the purity of operations is used to parallelize the execution plan, improving performance.

We completely eliminate the need for users to write value inspection code, by providing context-menu actions on variables to inspect their values (Sec. VI). In an educational context, where value inspection steps must be contained explicitly in the notebook, we propose to tag respective cells, so that they can be eliminated when running a pipeline (Sec. V).

Overall, we expect that our approach will significantly speed up DS pipeline development, by (1) avoiding bugs resulting from access to stale notebook state, (2) rerunning only the code that must be rerun to update stale state, and (3) eliminating the need to write and understand inspection code. We will verify this claim in a usability study once the prototype for Safe-DS is complete.

Acknowledgments

This work was partially funded by the Federal Ministry of Education and Research (BMBF), Germany under the Simple-ML project (grant 01IS18054). We thank the reviewers for their numerous constructive comments and suggestions.