diff --git a/python/ray/data/dataset.py b/python/ray/data/dataset.py
index 929373d94b008..43268a5bd8fb5 100644
--- a/python/ray/data/dataset.py
+++ b/python/ray/data/dataset.py
@@ -713,6 +713,7 @@ def add_column(
 
             >>> import ray
             >>> import pyarrow as pa
+            >>> import pyarrow.compute as pc
             >>> ds = ray.data.range(100)
             >>> ds.schema()
             Column  Type
@@ -721,7 +722,7 @@ def add_column(
 
             Add a new column equal to ``id * 2``.
 
-            >>> ds.add_column("new_id", lambda x: pa.compute.multiply(x["id"], 2)).schema()
+            >>> ds.add_column("new_id", lambda x: pc.multiply(x["id"], 2)).schema()
             Column  Type
             ------  ----
             id      int64