Skip to content

Closes #5228: remove type ignore from factorize in extension module #5229

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
ajpotts merged 3 commits into from
Jan 16, 2026
Merged

Closes #5228: remove type ignore from factorize in extension module #5229

ajpotts merged 3 commits into from
Jan 16, 2026
+129 −157

Conversation

@ajpotts
Copy link
Contributor

@ajpotts ajpotts commented Jan 5, 2026 •
edited
Loading

PR: Align ExtensionArray factorize and argsort with pandas expectations

Summary

Align pandas ExtensionArray behavior with pandas expectations by returning NumPy arrays
(not Arkouda arrays) for factorize codes and argsort indices, while keeping all grouping
and sorting computation server-side in Arkouda.

This improves pandas compatibility, simplifies downstream pandas internals
(e.g. groupby, take, iloc), and clarifies API semantics.

Key changes

ArkoudaExtensionArray.factorize

  • Return type updated to: 
    Tuple[np.ndarray, ArkoudaExtensionArray]
  • codes are now returned as a NumPy array of dtype np.intp, as expected by pandas.
  • uniques are returned as an ExtensionArray of the same type as self.
  • Removed sort argument:
    • Factorization now consistently uses first-appearance order, matching pandas’ default behavior.
  • Clarified missing-value handling:
    • Only floating dtypes treat NaN as missing.
    • use_na_sentinel controls whether missing values map to -1 or len(uniques).

ArkoudaExtensionArray.argsort

  • Return type changed from Arkouda pdarray to NumPy ndarray[np.intp].
  • Sorting computation remains server-side; only the permutation indices are materialized client-side.
  • na_position is now accepted via **kwargs for pandas compatibility.
  • Updated docstring to reflect pandas ExtensionArray contract more precisely.

Tests

  • Updated tests to assert against NumPy arrays instead of Arkouda arrays for:
    • factorize codes
    • argsort indices
  • Removed sort-parameter test cases and adjusted expectations to first-appearance semantics.
  • Added explicit NumPy assertions (numpy.testing.assert_equal) for clarity and correctness.

Motivation

Pandas internals expect:

  • factorize → NumPy integer codes
  • argsort → NumPy permutation indices

Returning Arkouda arrays in these paths caused unnecessary friction and divergence from the
pandas ExtensionArray contract. This PR preserves Arkouda’s distributed execution model
while presenting pandas-native results at the API boundary.

Closes #5228: remove type ignore from factorize in extension module

@ajpotts ajpotts force-pushed the 5228_remove_type_ignore_from_factorize_in_extension_module branch from 74ee5f0 to 937d368 Compare January 5, 2026 21:30
@ajpotts ajpotts marked this pull request as ready for review January 6, 2026 00:35
1RyanK
1RyanK approved these changes Jan 14, 2026
View reviewed changes
Copy link
Collaborator

@1RyanK 1RyanK left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good!

@ajpotts ajpotts force-pushed the 5228_remove_type_ignore_from_factorize_in_extension_module branch from 937d368 to bfdfb5a Compare January 15, 2026 10:42
@codecov
Copy link

codecov bot commented Jan 15, 2026 •
edited
Loading

Codecov Report

✅ All modified and coverable lines are covered by tests.
⚠️ Please upload report for BASE (main@983fc7a). Learn more about missing BASE report.

Additional details and impacted files 
@@           Coverage Diff            @@
##             main     #5229   +/-   ##
========================================
  Coverage        ?   100.00%           
========================================
  Files           ?         4           
  Lines           ?        63           
  Branches        ?         0           
========================================
  Hits            ?        63           
  Misses          ?         0           
  Partials        ?         0

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

jaketrookman
jaketrookman approved these changes Jan 15, 2026
View reviewed changes
Copy link
Contributor

@jaketrookman jaketrookman left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good

@ajpotts ajpotts enabled auto-merge January 15, 2026 19:20
drculhane
drculhane approved these changes Jan 15, 2026
View reviewed changes
Copy link
Contributor

@drculhane drculhane left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Certainly seems to do what's advertised. Nice work.

@ajpotts ajpotts added this pull request to the merge queue Jan 16, 2026
Merged via the queue into Bears-R-Us:main with commit 538fb39 Jan 16, 2026
24 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@jaketrookman jaketrookman jaketrookman approved these changes

@1RyanK 1RyanK 1RyanK approved these changes

@drculhane drculhane drculhane approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

Implement Pandas ExtensionArray API

Development

Successfully merging this pull request may close these issues.

remove type: ignore from factorize in extension module

4 participants

@ajpotts @jaketrookman @1RyanK @drculhane