fix: row_id range fix for index training on gpu

[jerryyifei]

The original sampling logic may create row_id range that is out of total_record length.
This diff makes sure the row_ids are within the range

Fixes #1662

[github-actions]

ACTION NEEDED

Lance follows the Conventional Commits specification for release automation.

The PR title and description are used as the merge commit message. Please update your PR title and description to match the specification.

For details on the error please inspect the "PR Title Check" action.

[rok]

Thanks for doing this @jerryyifei !
I've left a couple of comments.

[rok]

Why change here?

[rok]

Ditto

[wjones127]

Thanks for making a PR! I think there's a simpler change that would do a better job of preserving the probability distribution.

[wjones127]

Instead of altering the outputs of sampling, why not fix the inputs? I think all you have to do is change chunk_size in the last iteration to be the remaining size (total_records - i) here.

local_size = min(chunk_size, total_records - i)

[rok]

That would also produce a more uniform distribution in the last chunk.

[rok]

How about if we simplify like this?

Suggested change

	num_sampled += chunk_sample_size
	# If we are at the last chunk, we may not have enough records to sample.
	local_size = min(chunk_size, total_records - i)
	local_sample_size = min(chunk_sample_size, local_size)
	if local_sample_size < local_size:
	# Add more randomness within each chunk, if there is room.
	offset = i + np.random.randint(0, local_size - local_sample_size)
	else:
	offset = i
	# If we are at the last chunk, we may not have enough records to sample.
	local_sample_size = min(chunk_sample_size, total_records - i)
	num_sampled += local_sample_size
	offset = i + np.random.randint(0, local_sample_size)

[wjones127]

That looks like a totally different computation for offset? Offset needs to be constrained such that offset + local_sample_size <= local_size, otherwise we can get duplicate rows sampled.

[rok]

I'm thinking:

        local_size = min(chunk_size, total_records - i)
        local_sample_size = min(chunk_sample_size, local_size)

combining:

        local_sample_size = min(chunk_sample_size, min(chunk_size, total_records - i)) 

simplify:

        # because chunk_size > chunk_sample_size
        local_sample_size = min(chunk_sample_size, total_records - i)

local_sample_size <= local_size is always true because local_sample_size = min(chunk_sample_size, local_size) so we can just say offset = i + np.random.randint(0, local_sample_size). Which I find more readable.
I hope I didn't miss something obvious :).

[rok]

I think the change is good either way, just looking to get rid of an extra variable.

[wjones127]

I still don't understand why

offset = i + np.random.randint(0, local_sample_size)

is valid.

We are trying to find a random offset such that a window of size local_sample_size is inside of the window of size local_size. Hence the invariant I mentioned:

offset + local_sample_size <= local_size

 ◄────────────── local_size ──────────────►

 ┌──────┬────────────────┬────────────────┐
 │      │                │                │
 │i     │i + offset      │i + offset + lss│i + local_size
 │      │                │                │
 └──────┴────────────────┴────────────────┘

 ◄──┬───► ◄──────┬───────►
    │            │
 offset          │
         local_sample_size

[rok]

Oh I see! The offset would have to be:

offset = i + np.random.randint(0, max(0, chunk_size - local_sample_size))

This is fine though :)

[wjones127]

Thanks for your help @jerryyifei!

[jerryyifei]

Thanks to @wjones127 @rok for help merge! Sorry I was out for a trip and didn't reply this thread on time.