Repository navigation
Improve performance of Table argsort and Time get_sortable_arrays - #20150
Conversation
…orting of mixins Also add regression tests for multi-key sorting and warnings on 'kind' argument.
|
Thank you for your contribution to Astropy! 🌌 This checklist is meant to remind the package maintainers who will review this pull request of some common things to look for.
|
mhvk
left a comment
There was a problem hiding this comment.
@taldcroft - the crucial bits are great, and my comments are just a request to reconsider the warning added, and one little thing that can be removed.
| t = Table() | ||
| t["a"] = [2, 1, 3, 2, 3, 1] | ||
| a_struct = t["a"].copy() | ||
| with warnings.catch_warnings(): |
There was a problem hiding this comment.
This is the default when we run tests, so no need to add
(aside: can do catch_warnings(action="error"))
| The column name(s) to order the table by | ||
| kind : {'quicksort', 'mergesort', 'heapsort', 'stable'}, optional | ||
| Sorting algorithm used by ``numpy.argsort``. | ||
| Sorting algorithm used by `numpy.argsort`. This is only used |
There was a problem hiding this comment.
It is not a big deal, but my sense would be never to warn and here just write
Passed on to `numpy.argsort` if sorting is on a single column,
ignored for sorting on multiple columns (including mixin columns
that expand into multiple sortable arrays, like |Time|).
My logic is that for sorting a warning is really not that useful, and for people who, say write a package that wraps table sorting, and who just want stable sorts (which is not the default for argsort), it is nice to be able to give it and not worry about how many columns are going to be sorted.
Note that np.argsort is not that different; the kind parameter is only a request (and mergesort doesn't really exist as a separate option anymore):
Sorting algorithm. The default is 'quicksort'. Note that both 'stable'
and 'mergesort' use timsort under the covers and, in general, the
actual implementation will vary with data type. The 'mergesort' option
is retained for backwards compatibility.
| kwargs = {"kind": kind} if kind else {} | ||
| idx = np.argsort(sortable_arrays[0], **kwargs) | ||
| else: | ||
| if kind is not None and kind not in ("stable", "mergesort"): |
There was a problem hiding this comment.
So, I'd remove the warning here (which obviously simplifies the tests too).
| jd1, jd2 = parent.jd1, parent.jd2 | ||
| jd_approx = jd1 + jd2 | ||
| jd_remainder = (jd1 - jd_approx) + jd2 | ||
| return [jd_approx, jd_remainder] |
There was a problem hiding this comment.
I really don't know why we didn't do this before!
|
@mhvk - thanks for the review. I implemented your suggestion to remove the warning, excellent idea. |
|
FYI the |
|
When did you last rebase? Might help. |
|
I see you have enabled auto-merge but it is not the CI blocking the merge. This PR needs at least one approval. FYI. |
Description
This PR use
col.info.get_sortable_arrays()inTable.argsort()for efficient mixin sorting. This fixes the first bullet point of #14942.Rationale
Table.argsort()(and henceTable.sort()) currently sorts by converting the key columns to a NumPy array viaTable.as_array(). For a single plain key column this is just the column data, but for multiple key columns it builds a structured array over all of them, and if any key is a mixin column (like~astropy.time.Time) that can't be represented as native structured-array dtype, it gets stored as a slow Python object array. Every comparison during the sort then falls back to Python-level object comparisons instead of vectorized NumPy operations, which is extremely slow — see the benchmark below, where combining aTimecolumn with a plain int column is over 1000x slower than sorting on the int column alone.get_sortable_arrays()already exists as aninfomethod (used internally by table joins) that returns one or more plain NumPy arrays which can be lexically sorted to reproduce the sort order of a column, without needing an object array. For a normalColumnthis is just[self.data]. For aTimecolumn it is two float arrays (jd1/jd2split into an approximate value and a remainder) that fully capture time ordering without loss of precision.This PR reworks
Table.argsort()to callget_sortable_arrays()on each key column and flatten the results into a single list of arrays. If there is only one sortable array overall (e.g. a single plain-column key),np.argsortis used directly, preserving support for thekindargument. Otherwise (multiple keys, and/or a mixin column that expands into multiple arrays)np.lexsortis used, which is always stable and does not acceptkind; in that case a non-default, non-stablekindis ignored and triggers anAstropyUserWarning.If a column's
get_sortable_arrays()raisesNotImplementedError(e.g.SkyCoord, which has no well-defined sort order),Table.argsort()now raisesTypeErrorrather than falling through to a slow/incorrect object array comparison.As part of this work,
Time.info.get_sortable_arrays()was also optimized: it previously computed the remainder term by constructing a newTimeobject and subtracting, which is roughly as expensive as computing a fullTime.argsort(). It now uses the same cheapjd1/jd2float arithmetic thatTime.argsort()already uses internally, avoiding that redundant object construction.Changes
astropy/table/table.py:Table.argsort()rewritten to usecol.info.get_sortable_arrays()per key column, dispatching tonp.argsortfor a single sortable array ornp.lexsortotherwise.astropy/time/core.py:TimeInfo.get_sortable_arrays()optimized to avoid constructing an intermediateTimeobject.astropy/table/tests/test_table.py: updatedtest_sort_kindfor the newkindwarning behavior on multi-key sorts, and addedtest_sort_kind_single_keycovering the single-sortable-array code path for allnumpy.argsortkindvalues.astropy/table/tests/test_operations.py: added regression tests for multi-key sorts involving aTimecolumn, thekindwarning on multi-key sorts, and theTypeErrorraised for non-sortable mixin columns (e.g.SkyCoord).docs/changes/table/20149.perf.rst: changelog entry.Backwards compatibility notes
get_sortable_arrays()(raisesNotImplementedError, e.g.SkyCoord) now raisesTypeErrorinstead of attempting (and likely failing or behaving incorrectly on) an object-array sort.Performance
Table.argsort()performance:mainvstable-unified-sortBenchmark script: perf_test_argsort_time.py
N = 100,000 rows, 5 repeats per case, median times reported.
Results
maintable-unified-sorttimeonly["time", "int"]["int1", "int2"]["int1"]onlyKey findings
get_sortable_arrays()path avoids building a structured array viaas_array().Timemixin column with a plain column onmain: sorting["time", "int"]forcesas_array(names=keys), which converts theTimecolumn into a structured array of Python objects, making every comparison during the sort go through slow object comparisons — over 1000x slower than this branch's approach of flatteningget_sortable_arrays()and usingnp.lexsortdirectly on numeric arrays.Timecombined with any other sort key were essentially unusable at scale onmain.AI disclosure
Essentially all of this PR was generated using GitHub CoPilot with Claude Sonnet 5. I reviewed everything carefully and understand all the changes.