• Tests added / passed
• Passes black dask / flake8 dask

• rebased ≈2021-05-20T01:00Z
• description below updated 2021-03-23
• Dask Summit lightning talk slides: bit.ly/dask-summit-iloc

DataFrame.iloc is generally unsupported, since it requires knowing the number of elements in each partition. DataFrames usually know each partition's index bounds (divisions), which allows implementing loc, but the partition sizes are never computed or saved.

In some cases, it is possible to infer, store, and propagate partition sizes at graph-construction time, so I make a best effort to do so here, and implement iloc for DataFrames that have the corresponding partition_sizes member set.

Many DataFrame-creation methods (converting from Pandas, chunked file formats, or Dask Arrays w/ known chunks) can seed partition_sizes, and many transformations can propagate existing partition_sizes, so a large proportion of DataFrames in the wild may be able to benefit from these new members.

Reviewing

I've broken the change into 12 sequential chunks of 5-15 fine-grained commits. Basic CI checks pass on every chunk, and they can in principle be reviewed and merged sequentially:

• p0..p1 (squashed: q0..q1)
  • add partition_sizes to DataFrame constructors and state
  • add preserve_partition_sizes flag to map_partitions that call-sites can set to True during "narrow" map steps, to ensure partition_sizes are propagated
  • add partition-size checking to assert_eq, similar to existing check_divisions functionality; also allow for more configurable verifications of both divisions and partition_sizes
• p1..p2 (q1..q2)
  • support from_pandas (results in many assert_eq tests starting to have partition_sizes values for one of the _Frames being checked)
  • support len backed by partition_sizes (when present)
  • is_scalar_for_elemwise tweak
• p2..p3 (q2..q3): from_dask_array, map, __getitem__(str|scalar)
• p3..p4 (q3..q4): rename, copy, some repartition paths, loc, some parquet support
• p4..p5 (q4..q5): fillna, align, more repartition, head, tail, mode, …
• p5..p6 (q5..q6): merge_asof, isna, get/set index, …
• p6..p7 (q6..q7): values, to_frame
• p7..p8 (q7..q8): apply, reset_index, binary methods
• p8..p9 (q8..q9): assign, stats/agg methods
• p9..p10 (q9..q10): __getitem__(list/ndarray), consolidate Series and DataFrame __getitem__ into _Frame, Array.__getitem__(Series), preserve delegation of Array ops to _Frame
• p10..p11 (q10..q11): _Frame.iloc, tests 🎉
• p11..ps (q11..qs): repartition_sizes using iloc, use repartition in to_dask_array

Notes:

• Branches starting with q are squashed versions of each corresponding p chunk (the tree at qX is the same as the tree at pX), in case that's easier to review (e.g. q0..qs is all the squashed chunks).
• The starting point, p0, is an octo-merge of open PRs I've factored out of this one (currently Improve handling of duplicate divisions (i.e. empty partitions) #7214, comment fixes #7215, Ensure natural sort order in parquet part paths #7249, and dead code removal / fixes #7388). q0 (squashed-chunk base) is the same commit as p0.
• Each chunk's final commit passes "CI Additional" checks (we disabled the full "CI" GHA in celsiustx/dask a while back, but probably worth re-enabling now so that we can get better confidence without spamming dask/dask with PRs)

One path forward would be for me to file separate PRs for the first chunk or couple chunks, to begin a real review process on them. Suggestions welcome.

Rebasing Notes

• I rebase p0 by rebasing its constituent branches, then using a git-octomerge helper:

  for b in nits divisions parquet; do gco $b && grb; done && gco main && gom nits dcr divisions parquet

• Then I rebase the rest of the p line using git-rebase-sequence:

  grq p0 <previous p0> p1 p2 p3 p4 p5 p6 p7 p8 p9 p10 p11 ps
  

• I update the "q*" (squashed-chunk) branches from the "p" branches using this script:

  gsqsq p0 $(for i in `seq 1 11` s; do printf "p$i:q$i "; done; echo)
  

Progress

I've factored out #7214, #7215, #7249 #7388 of this PR.

#7396 is a prefix of this PR, and can in principle be reviewed+merged first.

Implementation Notes

A few pointers to impt new code blocks (I'll try to keep these updated as I update the code):

• iloc implementation (IlocIndexer.__getitem__) (could use some factoring!)
• iloc tests

Identifying _Frames with "shape contains a Delayed" (if any(is_dask_collection(x) for x in shape))

In several places, "shape contains a Delayed" is used to detect whether something is a DataFrame or Series. That relationship was a bit cryptic/indirect, and is less straightforward now (as many _Frames now know their shapes), so I tweaked it in various places.

Here is one example, in is_scalar_for_elemwise.

Another example is: most array binary ops were previously returning NotImplemented due to identifying arguments as _Frames in this manner (which allowed _Frame's version of the ops to be used). Those ops now explicitly import _Frame and do an isinstance test. Importing dask.dataframe from within dask.array seems like a smell, but it is apparently already done (always in a local scope, to avoid circular imports at the top level), and I'm not sure of another way to handle the situation. code

# TODO: partition_sizes

I left a # TODO: partition_sizes next to DataFrame transformations where I wasn't sure whether partition_sizes can be propagated (or haven't gotten around to implementing it; this commit contains many examples); commonly this occurs on map_partitions and new_dd_object calls.

Here's a command for finding them:

git grep 'TODO.*partition.sizes'

Setting df.partition_sizes = None in tests

Most DataFrame tests construct DataFrames via from_pandas, which now populates partition_sizes. This can result in tests failing to cover non-partition-size-aware code paths (which to date have been the default/only code paths), and makes it easy to regress or miss coverage for that functionality.

In a few tests, I explicitly set partition_sizes = None on an already-constructed DataFrame (example). I'm not sure if disabling partition_sizes like this should be a kwarg on from_pandas, or more tests should be parameterized to explicitly test the partition-sizeful and partition-sizeless paths, or what. Open to discussion.

TODOs

• more documentation
• factor iloc implementation (IlocIndexer.__getitem__)
• testing:
  • more coverage likely needed
  • performance: check_divisions and check_partition_sizes in assert_eq each compute DDFs' partitions; they can be combined for efficiency
  • perhaps allow those check_* params to take string values that indicate that each DDF's divisions/partition_sizes should be verified to be internally consistent (or None), but are allowed/expected to not equal each other (e.g. when testing repartitioning, or when division/partition_size info is expected to be lost during specific transformations)
• some TODOs in the code should probably be resolved before merging; closer look at those warranted (git grep 'TODO.*partition.sizes')
• suspected backwards overlap check should be factored out as own PR, specifically tested I think this is done
• verify parquet num-rows statistics correct (esp in presence of multiple row-groups; I'm guessing tests already cover this, but I played a bit of whack-a-mole getting them all passing across {fastparquet,pyarrow}) I think this is done
• double-check long-tail of aggregation functions where partition_sizes should be [the number of output columns] (mean,var,sem are tested and believed correct, but others may not be) I think this is done

Follow-on work

• infer partition_sizes in more DDF construction methods (e.g. HDF5, other formats?); can be follow-on PR
• fix from_pandas wonky partitioning (partition_sizes is correct based on existing implementation, it just sometimes results in skewed chunks for no good reason); can be follow-on PR