Dask data loss with `.set_index`

.env

#VERSION=2020.12.0
VERSION=2021.1.1

docker-compose.yml

version: '3'

services:
  scheduler:
    image: daskdev/dask:${VERSION}
    command: dask-scheduler
    volumes:
      - ./test.py:/srv/test.py
  
  wrk:
    image: daskdev/dask:${VERSION}
    command: dask-worker tcp://scheduler:8786
    deploy:
      replicas: 2

reproduce_issue.sh

docker-compose up -d

docker-compose exec scheduler python /srv/test.py

test output

(base) ➜  docker-compose exec scheduler python /srv/test.py
10000
10000
5015

test.py

import uuid
import pandas as pd
import dask.dataframe as dd
from dask.distributed import Client

with Client('scheduler:8786') as client:
    test_ddf = dd.from_pandas(pd.DataFrame({
        'uuid': [str(uuid.uuid4()) for i in range(10000)],
    }), chunksize=100)

    print(len(test_ddf))
    print(len(test_ddf.set_index('uuid')))
    print(len(test_ddf.set_index('uuid', shuffle='disk')))