Use Databricks Assistant

Databricks Assistant is a context-aware AI assistant that is available in Databricks notebooks, SQL editor, and file editor. Databricks Assistant lets you query data using a conversational interface, making you more productive inside Databricks. You can describe your task in English and let the assistant generate Python code or SQL queries, explain complex code, and automatically fix errors. The assistant uses Unity Catalog metadata to understand your tables, columns, descriptions, and popular data assets across your company to provide personalized responses.

Databricks Assistant can perform the following tasks:

  • Generate code.

  • Debug code, including identifying and suggesting fixes for errors.

  • Transform and optimize code.

  • Explain code.

  • Help you find relevant information in the Databricks documentation.

Databricks Assistant is enabled by default. An admin can disable or enable Databricks Assistant for all workspaces in an account. If an admin has permitted workspace setting overrides, workspace admins can enable or disable Databricks Assistant for specific workspaces. For more details, see Enable or disable Databricks Assistant.

Enablement of the Databricks Assistant for your account is captured as an account event in your audit logs, see Account-level account events.

Additional information

The following articles contain additional information about using Databricks Assistant:

How to use Databricks Assistant

Databricks Assistant is available in notebooks, SQL editor, file editor, and dashboards. This document focuses on notebooks. Usage in the other products is similar. For information about using Databricks Assistant to create visualizations for dashboards, see Create visualizations with Databricks Assistant.

In a notebook, Databricks Assistant is available in the Assistant pane or inline in a code cell.

Use Databricks Assistant in a notebook cell

To use Databricks Assistant directly in a code cell, press Cmd + i on MacOS or Ctrl + i on Windows. A text box appears within the cell. You can type a question or comment in English and then press Enter (not Shift+Enter, as you would to run a cell) to have the assistant generate a response.

screenshot of inline assistant in use

You can also type / to display a list of common commands, such as /explain to explain the code in a cell, /fix to propose a fix if the code has an error, and /doc to suggest documentation comments for your code.

When you use /fix or /doc, the assistant displays proposed changes in a diff window. Select Accept to accept the proposed changes or Reject to keep the original code. If you accept the proposed code, the code does not automatically run. You can review the code before running it. If the generated code is not what you wanted, try again by adding more details or information to your comment. See Tips for using Databricks Assistant.

For code autocomplete, performance may be better using the Assistant pane than in a notebook cell.

To exit the Assistant in a notebook cell, click close assistant in the upper-right corner of the cell. The Assistant closes automatically if you Accept or Reject the code it generated.

Use the Assistant pane

To open the Assistant pane, click Databricks assistant icon in the left sidebar.

screenshot of assistant pane in use

Type questions in the text box at the bottom of the Assistant pane and press Enter or click enter assistant text at the right of the text box. The Assistant displays its answer. The following screenshot shows actions you can take after the Assistant has generated code in the Assistant pane.

icons at the top of the code box in the assistant pane

You can run the same query again to generate another answer. To do so, hover your cursor over the answer and click regenerate query icon.

To close the pane, click the icon again or click close assistant in the upper-right corner of the cell. You can expand the pane to full width by clicking open full width icon; click close full width icon to return the pane to default width.

icons at the top of the assistant pane

The Assistant pane keeps track of your conversations even if you close the pane or notebook. To clear previous conversations, click clear-assistant at the upper-right of the Assistant pane.

Debug code

To use Databricks Assistant to fix code, ask a question in the Assistant pane or click the Diagnose Error button that appears in the cell results when an error occurs.

assistant debugging example

Transform or optimize code

Databricks Assistant can help optimize your code to make it more efficient. For example, it can suggest ways to optimize slow or inefficient code such as replacing unnecessary loops with built-in functions, or it can convert pandas code to PySpark or Spark SQL.

Explain code

Databricks Assistant can provide detailed explanations of code snippets. In the prompt you provide, you can include terms like “be concise” or “explain code line-by-line” to request the level of detail for the explanation it provides. You can also ask Databricks Assistant to add comments to code.

Find relevant information in the Databricks documentation

Databricks Assistant can help answer questions based on Databricks documentation directly from the notebook editor.

answer question based on Databricks documentation

Tips for using Databricks Assistant

This section includes some general tips and best practices when using Databricks Assistant.

Databricks Assistant uses context to provide better answers

Databricks Assistant has accesss to table and column schema and metadata. This allows you to use natural language and generate more accurate queries. For example, if a table has a column State, you can ask Databricks Assistant to generate a list of users who live in Michigan.

Databricks Assistant uses the following context:

  • Code or queries in the current notebook cell or Databricks SQL editor tab.

  • Table and Column names and descriptions.

  • Previous prompt questions.

  • Favorite and active tables.

  • For the diagnose error feature, the stack trace from the error output.

When selecting columns from a DataFrame, you can get more accurate results by providing a starting query. For example, provide a statement like SELECT * FROM <table_name>. This allows Databricks Assistant to get the column names and not have to guess.

Since Databricks Assistant uses your conversation history to provide better and more accurate answers, you can ask Databricks Assistant to alter the output of a previous response without having to rewrite the entire prompt, as shown below.

Similar to the above example, you can use the Assistant’s chat history to iteratively clean, explore, filter, and slice DataFrames in the Assistant pane.

Automatically format code

To format your code to adhere to best practices, type “prettify” into the assistant text box or assistant pane.

format code example

Be specific

Because Databricks Assistant does not use row-level data, you can provide more detail to the prompts to get the most accurate answer. For example, specify what the data looks like or if column data type conversions are needed to perform an operation.

For example: ”convert this code from pandas to PySpark, including the code needed to convert the pandas DataFrame to a PySpark DataFrame and changing the data type of column churn from boolean to integer”.

Use Shift+Enter to add a new line in the chat text box

You can use Shift+Enter to add a new line in the Assistant chat text box. This makes it easy to format and organize your messages to Databricks Assistant.

Edit and run code in Databricks Assistant chat pane

You can run code in the Assistant pane to validate it or use it as a scratchpad. To run code, click run code icon in the upper-left corner of the code box in the Assistant pane.

run code in assistant pane

When you run code in the Assistant pane, output is displayed and the variables become usable in the notebook.

You can also edit the code that Databricks Assistant generates directly in the Assistant chat box before moving the code to the notebook.