Rework data loaders

[westonpace]

We have a basic data loader for pytorch. However, this data loader has several limitations and is not very intuitive. There are also other training libraries (e.g. Ray Train, TensorFlow, etc.) and it would be good to ensure a consistent experience across all of them. This issue captures a lot of feedback (and we welcome more!) that we have gathered over the years in hopes of motivating a new design.

Clearly document fork-safety

Spawn is considerably slower than fork. In addition, users have often missed our guidelines on fork safety in the past and ended up frustrated with a hung process. I think #3584 will help here but there is more we could do.

• Improve our documentation on fork safety (we should improve our documentation on data loading in general, see below)
• Make lance fork-safe. I think this would mostly involve resetting static state (e.g. RT, LOOP) after a fork. However, this is very difficult to guarantee in general and would likely be an on-going challenge.
• Document how to use forkserver. It is very possible to safely use forkserver in lance. One just needs to take care to ensure that lance is loaded after the forkserver is started. Improved documentation on data loading should demonstrate this.

Provide both map-based and iterator-based datasets

We currently only have an implementation of torch's iterator-based datasets. Several users have created their own map-based data loaders and these are useful in a variety of situations. However, properly implementing a map-based data loader can be tricky (e.g. need to understand row addresses and pin to version, should use slightly hidden __getitems__ method, not clear how to do shuffling with restarts, etc.)

Simplify filtering & shuffling

The filtering and shuffling in the current data loader is confusing. Pytorch's sampler is confusing. I think a much simpler model could be devised in Lance. From Lance perspective, these are the shuffling scenarios I am aware of and the approaches that should be taken.

Scenario	Approach
No shuffling, No filter	If there is no shuffling and there is no filter then an iterable dataset should be used for best performance
No shuffling, With filter	If there is no shuffling but there is a filter then we should decide between scan (iterable dataset) and take (map dataset) based on the filter selectivity (assuming there is an index on the filter we can do this cheaply)
Shard shuffling, No filter	If there is no filter and the user wants shard shuffling then an iterable dataset should be used. The shard size should be completely independent of the fragment size. There is no benefit to joining these two things and it muddies the waters. Offsets & limits should be used for sharding.
Shard shuffling, With filter	If the filter is not very selective then we could still use an iterable dataset and filter after partitioning. If the filter is highly selective then falling back to a map dataset is probably easiest. Can suggest materialized view or local storage if higher performance is desired.
Row shuffling, With/without filter	If row shuffling is desired then we should shuffle the (possibly filtered) row ids and run a map dataset on top of this. Consider materialized view or local storage if performance is not sufficient.

Checkpoints

Model training creates checkpoints. All of the above data loading scenarios should support creating and restarting from checkpoints. For example, when shard shuffling, we can give each shard a unique ID and save off which shards have been processed and resume the process at a later date.

This will require both documentation and features (e.g. being able to create a data loader from a partial state). The actual act of saving the checkpoint is potentially out of scope for Lance as the training library may have support for this already. However, we should clearly document how this is done in the major training libraries. Alternatively, we could offer a basic checkpoint storage mechanism.

Parallelize / batch conversion to tensors, provide more examples

The data loader currently has a to_tensor_function. However, in many cases, conversion to tensors can be a very expensive process. For example, we may be decoding images. Since the current function rows on a row-by-row basis and is not parallelized then it can easily become the bottleneck in data loading. We should make it possible to parallelize this process.

Better tracing / profiling

We should provide APIs to get basic performance information from the data loader so that it is easier to identify and fix bottlenecks. For example, after a run we could log:

• How much time was spent waiting for I/O, what was our effective I/O bandwidth?
• How much time was spent decoding and converting to tensor?
• How many rows per second did the data loader emit? How many bytes per second did it emit?

Create a complete guide for data loading

This can cover all of the above as well as give performance recommendations for several scenarios. For example, we could discuss the benefits of offline processing and creating materialized views (copies) of the data if there are many epochs to run.

[Jay-ju]

This is a very useful feature. #3641. In fact, it is also to solve part of this problem. We can take a look at the problem here together.

[jd-nuva]

Thanks for writing this up! It will be tremendously helpful for lance adoption since current API + documentation suggests they're many different paths to get the same training job done and it's largely left to the users to figure out and decide on the path, often not fully aware of which path would best align with lance's internal implementation for best efficiency and reliability.

Our training use case falls into Shard shuffling, With filter bucket where:

1. Each row is 4MB ~ 10MB+ in size since it's image/video data with various resolution buckets
2. Filter is necessary to ensure we filter training data within a batch to have the same resolution and tensor shape during training
3. Shuffling is must have for better training quality in case of multiple epochs, ideally shuffling should happen at reasonably small shard level between epochs, instead of repeatedly read from the same large fragment in the same order
4. Deterministic loading helps with checkpointing and fault tolerance so in between restarts, we're not consuming more or less data than expected that would skew the training data distribution
5. Multi-loader multi-host is highly preferable to work in distributed setting, i assume this is given if sampler / data loader interfaces are designed correctly.

[LawJarp-A]

Hi, is there active work going on here? I would like to contribute, I heavily used lance data-loaders in PyTorch and I would love to work on integrating it mainstream.
In fact, I don't know if you remember, I had a short call with @westonpace and Benjamin Lees, where I pointed out this exact thing. Would be happy to contribute.

[westonpace]

@LawJarp-A

I'm working on data loading this week actually 😆 but I'm working in LanceDB at the moment and not lance, because a lot of the work I'm doing is based around views, and those are better represented at a higher level. Here's what I've been doing and plan to do this week...

• I'm adding a new system column (_rowoffset) which represents the index of the row in the overall dataset. By reading and filtering on this column we get a more general equivalent of the LanceDataset::take method. (feat: allow reading the _rowoffset #4478)
• I added take_offsets and take_row_ids to lancedb's Table. (feat: add take_offsets and take_row_ids lancedb/lancedb#2584)
• I'm creating utilities for "permutation views". This gives us a lot of the behavior that we currently need to rely on the torch data loader for like splits, shuffling, filtering, etc. (feat: a utility for creating "permutation views" lancedb/lancedb#2552)
• I've got another PR coming soon which will provide the read path for the permutation views.
• I've got another PR coming soon which will add __getitems__ to lancedb's Table which will allow both tables and permutation views to be used in a torch data loader

[LawJarp-A]

This looks great @westonpace, feels like I'm late to the party xD

If there's anything I can pick up from your task list, let me know! I've mostly been using (and hacking on) Lance for large image datasets, so I've got some experience there that might be useful.

The current to_tensor functionality is something I can work on - in my use case, I saw that it doesn't handle complex data types well, plus like you mentioned, the processing is sequential and row-by-row rather than batched, so slow as well

Happy to start wherever would be most helpful! Full disclosure tbh, this is my first OS contribution to the lance community, but I did play around a lot with the lance codebase wrt to torch data utils.

[westonpace]

The current to_tensor functionality is something I can work on - in my use case, I saw that it doesn't handle complex data types well, plus like you mentioned, the processing is sequential and row-by-row rather than batched, so slow as well

@LawJarp-A

Oh, good point, I haven't really started to dabble in this at all but I agree with your findings. Are you interested in working with Rust or sticking with Python?

[LawJarp-A]

@westonpace I primarily work in Python, but I've been learning rust from the past year and would not mind working on that side as well.

[westonpace]

@LawJarp-A here's my thinking on tensor functions:

First, we probably don't need to reinvent the old LanceDataset's to_tensor. Instead we should use the pytorch data loader's collate_fn which serves the same purpose.

1. Arrow -> PyTorch (lancedb work needed)

We need a good cheap way to convert Arrow record batches into pytorch tensors. I added a very basic to_tensor function in my test in lancedb/lancedb#2596 using DLPack:

def tbl_to_tensor(tbl):
    def to_tensor(col: pa.ChunkedArray):
        if col.num_chunks > 1:
            raise Exception("Single batch was too large to fit into a one-chunk table")
        return torch.from_dlpack(col.chunk(0))

    return torch.stack([to_tensor(tbl.column(i)) for i in range(tbl.num_columns)])

This could be improved. It's better than the once in lance for performance but still has some flaws. It works if all the columns are floats or signed integers. It doesn't work if any of the return types are FSL (e.g. tensors). The error may a little confusing if one of the columns is a different type. We could probably improve that. I'm also not sure it supports all the different kinds of types that both Arrow and PyTorch support. Also, this function should be moved into lancedb and not part of a unit test. It should also be documented. PyTorch users can then use this as their collate_fn whenever loading from a lancedb table.

2. Binary -> Arrow (lance work needed)

Often, the to_tensor function is not just converting Arrow to PyTorch, but also converting binary data into tensors. For example, parsing images or videos. In this case I don't think a to_tensor or collate_fn is the right spot to do the conversion. For one thing, it is pytorch only and not usable in other frameworks.

If we could make those operations an SQL UDF then users could set them in the projection (the projection can be set when creating a permutation). This would take care of batching and threading for us.

This would mean building libraries of UDFs for image processing, video processing, etc. We wouldn't need to actually write image processing code. We should be able to use existing libraries like opencv or magick-rust. We would just need to wrap the functionality of those libraries in SQL UDFs. This would mean:

• Deciding which functions to turn into UDFs
• Writing code to convert from arrow binary arrays into whatever input those libraries expect
• Writing code to convert the output from those functions into Arrow data types (e.g. opencv's imread function returns a Mat which is a pixel raster but how do we get that pixel raster into an Arrow FixedSizeList tensor type? Ideally with zero copies if possible)

[LawJarp-A]

@westonpace
Keeping data processing in the SQL layer and tensor conversion at the framework level is much cleaner

Some key functions (easy wins) top of my head

Image Processing UDFs

• decode_image(bytes)
• resize_image(tensor, width, height)
• crop_image(tensor, x, y, width, height)
• normalize_image(tensor, mean, std)
• to_rgb(tensor)

We can extend it for video/text subsequently once we have the process in place I feel. Most of the use cases are image/text
And like you said we could use existing libraries

As for the FSL, should be straightforward to imp0lementy it to to_tensors, just need to detect when we have Fixed Size List columns and handle the reshape from flat arrays back to tensor dimensions

[LawJarp-A]

@westonpace Let me know if I can get started on anything :)
UDF functions seems pretty independent as of now