[Datasets] Re-enable Parquet sampling and add progress bar

[c21]

Signed-off-by: Cheng Su [email protected]

Why are these changes needed?

This PR is to re-enable Parquet files sampling after testing multiple time in xgboost 100GB nightly test - https://github.com/ray-project/ray/blob/master/release/air_tests/air_benchmarks/workloads/xgboost_benchmark.py .

Verified no OOM and memory usage increase during training and tuning:

Also add a progress bar Parquet Sample Progress for the sampling to give us more information in case sampling takes a long time (which is unexpected). Tested on sampling 10 files from 100GB dataset, and it takes 3 seconds to finish:

(base) ray:~/workspace-project-cheng-debug% TEST_OUTPUT_JSON=/tmp/release_test_out.json python release/air_tests/air_benchmarks/workloads/xgboost_benchmark.py
Metadata Fetch Progress:   0%|          | 0/100 [00:00<?, ?it/s]
Metadata Fetch Progress:   1%|          | 1/100 [00:02<04:37,  2.80s/it]
Metadata Fetch Progress:   3%|▎         | 3/100 [00:02<01:17,  1.25it/s]
Metadata Fetch Progress:   6%|▌         | 6/100 [00:03<00:31,  3.02it/s]
Metadata Fetch Progress:   9%|▉         | 9/100 [00:03<00:18,  4.96it/s]
Metadata Fetch Progress:  11%|█         | 11/100 [00:03<00:18,  4.69it/s]
Metadata Fetch Progress:  13%|█▎        | 13/100 [00:03<00:15,  5.77it/s]
Metadata Fetch Progress:  15%|█▌        | 15/100 [00:04<00:14,  6.01it/s]
Metadata Fetch Progress:  19%|█▉        | 19/100 [00:04<00:11,  7.25it/s]
Metadata Fetch Progress:  21%|██        | 21/100 [00:04<00:10,  7.21it/s]
Metadata Fetch Progress:  24%|██▍       | 24/100 [00:05<00:07,  9.67it/s]
Metadata Fetch Progress:  26%|██▌       | 26/100 [00:05<00:07,  9.69it/s]
Metadata Fetch Progress:  28%|██▊       | 28/100 [00:05<00:07, 10.15it/s]
Metadata Fetch Progress:  30%|███       | 30/100 [00:05<00:05, 11.69it/s]
Metadata Fetch Progress:  32%|███▏      | 32/100 [00:05<00:07,  9.57it/s]
Metadata Fetch Progress:  34%|███▍      | 34/100 [00:05<00:05, 11.21it/s]
Metadata Fetch Progress:  36%|███▌      | 36/100 [00:06<00:06, 10.05it/s]
Metadata Fetch Progress:  38%|███▊      | 38/100 [00:06<00:06,  9.76it/s]
Metadata Fetch Progress:  42%|████▏     | 42/100 [00:06<00:04, 13.73it/s]
Metadata Fetch Progress:  44%|████▍     | 44/100 [00:06<00:03, 14.70it/s]
Metadata Fetch Progress:  49%|████▉     | 49/100 [00:06<00:02, 19.39it/s]
Metadata Fetch Progress:  52%|█████▏    | 52/100 [00:06<00:02, 19.46it/s]
Metadata Fetch Progress:  56%|█████▌    | 56/100 [00:07<00:02, 20.82it/s]
Metadata Fetch Progress:  59%|█████▉    | 59/100 [00:07<00:02, 13.80it/s]
Metadata Fetch Progress:  61%|██████    | 61/100 [00:07<00:02, 13.57it/s]
Metadata Fetch Progress:  63%|██████▎   | 63/100 [00:07<00:02, 13.69it/s]
Metadata Fetch Progress:  65%|██████▌   | 65/100 [00:08<00:03, 10.29it/s]
Metadata Fetch Progress:  67%|██████▋   | 67/100 [00:08<00:02, 11.48it/s]
Metadata Fetch Progress:  69%|██████▉   | 69/100 [00:08<00:03,  8.60it/s]
Metadata Fetch Progress:  71%|███████   | 71/100 [00:09<00:03,  7.61it/s]
Metadata Fetch Progress:  74%|███████▍  | 74/100 [00:09<00:02,  8.78it/s]
Metadata Fetch Progress:  78%|███████▊  | 78/100 [00:09<00:01, 12.49it/s]
Metadata Fetch Progress:  80%|████████  | 80/100 [00:09<00:01, 13.15it/s]
Metadata Fetch Progress:  82%|████████▏ | 82/100 [00:09<00:01, 12.96it/s]
Metadata Fetch Progress:  84%|████████▍ | 84/100 [00:10<00:02,  7.87it/s]
Metadata Fetch Progress:  86%|████████▌ | 86/100 [00:10<00:01,  7.54it/s]
Metadata Fetch Progress:  88%|████████▊ | 88/100 [00:10<00:01,  8.22it/s]
Metadata Fetch Progress:  90%|█████████ | 90/100 [00:10<00:01,  9.22it/s]
Metadata Fetch Progress:  92%|█████████▏| 92/100 [00:11<00:01,  6.61it/s]
Metadata Fetch Progress:  93%|█████████▎| 93/100 [00:13<00:03,  1.95it/s]
Metadata Fetch Progress:  96%|█████████▌| 96/100 [00:13<00:01,  3.23it/s]
Metadata Fetch Progress:  99%|█████████▉| 99/100 [00:14<00:00,  4.26it/s]
Metadata Fetch Progress: 100%|██████████| 100/100 [00:14<00:00,  7.04it/s]
Parquet Sample Progress:   0%|          | 0/10 [00:00<?, ?it/s]
Parquet Sample Progress:  10%|█         | 1/10 [00:00<00:05,  1.63it/s]
Parquet Sample Progress:  30%|███       | 3/10 [00:00<00:01,  4.20it/s]
Parquet Sample Progress:  40%|████      | 4/10 [00:00<00:01,  5.17it/s]
Parquet Sample Progress:  60%|██████    | 6/10 [00:01<00:00,  6.53it/s]
Parquet Sample Progress:  70%|███████   | 7/10 [00:02<00:01,  2.17it/s]
Parquet Sample Progress:  90%|█████████ | 9/10 [00:02<00:00,  3.47it/s]
Parquet Sample Progress: 100%|██████████| 10/10 [00:03<00:00,  3.26it/s]

Related issue number

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[c21]

Checked the failed tests and looks like irrelevant to this PR.