[ENH]: replace get_* methods on Arrow blocks with get_range()

[codetheweb]

Description of changes

Replaces specialized methods like get_gt and get_lt with a single get_range() method that behaves similarly to the std BTreeMap::range() method. This reduces complexity/repetition and also enables queries that are bounded in both directions.

Test plan

How are these changes tested?

• Tests pass locally with pytest for python, yarn test for js, cargo test for rust

Documentation Changes

Are all docstrings for user-facing APIs updated if required? Do we need to make documentation changes in the docs repository?

[codetheweb]

• [ENH]: make get_range() a lazy Stream #3009
• [ENH]: replace .get_* methods on blockfile API with .get_range() #2936
• [ENH]: replace get_* methods on memory blockfile impl with get_range() #2935
• [ENH]: replace get_* methods on Arrow blocks with get_range() #2934 👈
• [ENH]: replace get_block_ids_* with get_block_ids_range() in SparseIndex #2921
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

Join @codetheweb and the rest of your teammates on Graphite

[github-actions]

Reviewer Checklist

Please leverage this checklist to ensure your code review is thorough before approving

Testing, Bugs, Errors, Logs, Documentation

• Can you think of any use case in which the code does not behave as intended? Have they been tested?
• Can you think of any inputs or external events that could break the code? Is user input validated and safe? Have they been tested?
• If appropriate, are there adequate property based tests?
• If appropriate, are there adequate unit tests?
• Should any logging, debugging, tracing information be added or removed?
• Are error messages user-friendly?
• Have all documentation changes needed been made?
• Have all non-obvious changes been commented?

System Compatibility

• Are there any potential impacts on other parts of the system or backward compatibility?
• Does this change intersect with any items on our roadmap, and if so, is there a plan for fitting them together?

Quality

• Is this code of a unexpectedly high quality (Readability, Modularity, Intuitiveness)

[codetheweb]

I think it should be possible to remove the Clone bound by allowing borrowed ranges but I'm not entirely sure how. Methods like block.get_range() should accept both borrowed and owned ranges for ease of use. (E.x. I don't want to require users to do block.get_range(&(0..10)))

[Sicheng-Pan]

Looks good to me

[Sicheng-Pan]

Glad we can maintain one less binary search impl in the codebase now

[Sicheng-Pan]

Sometime key may not be required to derive an order (for example, when searching for start/end of prefix range). Since this is going to be the hot path in the system, I'm wondering if it's worth it to make key evaluation lazy?

[codetheweb]

I thought about that but didn't see a clean way to do it. Do you have ideas?

[Sicheng-Pan]

Maybe we can ask the comparator to take in a key handle:

• Before: F: FnMut((&'me str, K)) -> Ordering
• After: F: FnMut(&'me str, C) -> Ordering, C: Fn() -> &'me K

And we pass in || K::get(...) to the comparator.

Alternatively, ask the caller to pass in two comparators, one for prefix and the other for key. The second comparator is invoked only when the first yields Ordering::Equal (basically using a then_with)

Maybe there are better ways to do this, but this is what I have in mind right now.

[codetheweb]

yeah, that works

[codetheweb]

this turns out to be somewhat difficult:

F: FnMut(&'me str, C) -> Ordering, C: Fn() -> &'me K:

Can't work without Boxing the key fetch callback.

Alternatively, ask the caller to pass in two comparators, one for prefix and the other for key. The second comparator is invoked only when the first yields Ordering::Equal (basically using a then_with)

Makes life annoying for callers because the key cmp function parameter must be optional. But if None is provided, Rust complains that that the parameter type is unknown.

---

Are you ok leaving as-is for now and revisiting if it shows up in flamegraphs?

[Sicheng-Pan]

I'm totally fine with the current impl unless there is evidence showing that fetching the key will significantly slow us down

[Sicheng-Pan]

I'm thinking if we should decompose this into two functions for better code clarity:

• fn binary_search_prefix_key<K>(&self, prefix: &str, key: &K) -> Result<usize, usize> where Ok(i) means the (prefix, key) combination is found at index i and Err(i) means it is not found but can be inserted at i to maintain order
• fn find_smallest_index_with_prefix(&self, prefix: &str) -> Result<usize, usize>, where Ok(i) means the prefix exists and starts at i and Err(i) means it is not found but can be inserted at i with a key value to maintain order

[codetheweb]

fn find_smallest_index_with_prefix(&self, prefix: &str) -> Result<usize, usize>, where Ok(i) means the prefix exists and starts at i and Err(i) means it is not found but can be inserted at i with a key value to maintain order

I think this has to return an Option, you can't infer an insert location from only a prefix

[Sicheng-Pan]

I assume in the case the where the prefix does not exist, you can insert the prefix with any key value at that location without disturbing the order

[Sicheng-Pan]

But an Option should be sufficient too

[codetheweb]

🤦 yes you're right
I'll leave as an Option for now so find_smallest and find_largest are similar

[sanketkedia]

Nice! Curious, why does f have to be FnMut as opposed to Fn. Can't it work on immutable references?

[Sicheng-Pan]

I believe that's based on the standard library implementation. Probably that leaves more flexibility for what can be passed in as a comparator function.

[sanketkedia]

this looks much cleaner. Thank you for cleaning it up!

[codetheweb]

Merge activity

• Nov 4, 1:28 PM EST: A user started a stack merge that includes this pull request via Graphite.
• Nov 4, 1:32 PM EST: Graphite rebased this pull request as part of a merge.
• Nov 4, 1:33 PM EST: A user merged this pull request with Graphite.