[data] Optimize sample_boundaries in SortTaskSpec (#39581)

Optimize sample_boundaries in SortTaskSpec to call numpy.quantile once to get all boundaries for each column.

This is much faster than the old impl when num_reducers is large (eg. 5000), because each time numpy.quantile is called it actually sorts the array and the old impl calls it num_reducers times for each column.

Signed-off-by: z4y1b2 <88138737+z4y1b2@users.noreply.github.com>

python/ray/data/_internal/planner/exchange/sort_task_spec.py

@@ -112,10 +112,11 @@ def sample_boundaries(
         # Ignore the 1st item as it's not required for the boundary
         for k, v in sample_dict.items():
             sorted_v = v[indices]
-            sample_dict[k] = [
-                np.quantile(sorted_v, q, interpolation="nearest")
-                for q in np.linspace(0, 1, num_reducers)
-            ][1:]
+            sample_dict[k] = list(
+                np.quantile(
+                    sorted_v, np.linspace(0, 1, num_reducers), interpolation="nearest"
+                )[1:]
+            )
         # Return the list of boundaries as tuples
         # of a form (col1_value, col2_value, ...)
         return [