remove Dataset.num_blocks() usages (#307)

scottjlee · web-flow · commit 5a840af05d48 · 2024-02-14T17:01:22.000-08:00
Signed-off-by: Scott Lee &lt;sjl@anyscale.com&gt;
diff --git a/xgboost_ray/data_sources/ray_dataset.py b/xgboost_ray/data_sources/ray_dataset.py
@@ -107,4 +107,4 @@ def get_n(data: "ray.data.dataset.Dataset"):
         """
         Return number of distributed blocks.
         """
-        return data.num_blocks()
+        return data._plan.initial_num_blocks()
