feat: general compression for value page buffer

[niyue]

This PR introduces general compression for value page buffers, starting with zstd, to reduce the on-disk size of all types of value arrays.

Here are the key details:

1. After some code exploration, I implemented this as a buffer encoder for ValueEncoder instead of as an independent physical encoder. Please let me know if this approach is suitable.
2. Enhancements to ValuePageScheduler.schedule_ranges allow it to read the entire buffer range for compressed buffers. To support this, I added a new buffer_size metadata to several buffer structs, populating these variables using metadata from Lance.
3. Currently, only zstd compression with the default level (level 3) is implemented. If this approach is deemed suitable, we can consider adding more general compression methods in the future.
4. In a specific test case, the original Lance file was 13MB. After applying zstd compression, its size was reduced to 7.4MB.

[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.

[niyue]

I am unsure if it is appropriate to handle decompression in the scheduler. However, since the compressed bytes read from disk may be reused multiple times, decompression should not occur in decode_into for optimal performance. Please advise on this. Thank you.

[westonpace]

Yeah, this is a bit tricky, I see your concern. Unfortunately, we really shouldn't be doing the decompression here in the scheduler. The scheduler is single threaded and by doing compression here we will fall behind in scheduling I/O.

Something simple, that should work, is to wrap the compression task in a Arc<Mutex<...>> that each of the decode tasks that need that range share. Something roughly like...

struct ValuePageDecoder {
    bytes_per_value: u64,
    compressed_data: Vec<Bytes>,
    uncompressed_data: Arc<Mutex<Option<Bytes>>>,
}

impl ValuePageDecoder {

  // Called from decode_into
  fn get_uncompressed_bytes(&self) -> Result<Bytes> {
    let uncompressed_bytes = self.uncompressed_data.lock().unwrap();
    if uncompressed_bytes.is_none() {
      *uncompressed_bytes = self.decompress()?;
    }
    Ok(uncompressed_bytes.clone())
  }

}

This is not ideal because it means the decoder threads will be blocked for a bit but it should be good enough to start with.

A more complicated solution might be to have a dedicated thread pool for decompression. Once I/O is available we start to decompress it in this thread pool. I think we could do this as part of the I/O scheduler. So we can have arguments to submit_request which explain that the data is compressed and includes the parameters needed to decompress it. This would mean that compression is not just an "encoding" but something more fundamental but this is probably ok. We probably also want to treat encryption similarly. It also means we lose a bit of thread locality but I think that is inevitable. Decompression / decryption is costly enough that we want to parallelize across columns.

[niyue]

Something simple, that should work, is to wrap the compression task in a Arc<Mutex<...>> that each of the decode tasks that need that range share

This is not ideal because it means the decoder threads will be blocked for a bit but it should be good enough to start with.

Got it. I've considered this approach and its potential downsides. If it is deemed acceptable, I will revise the implementation accordingly.

[niyue]

I made the change accordingly, and uncompressed_bytes.clone() seems to copy the data, so I use Arc to avoid the copy (but the lock is still hold during decode_buffer), could you please check it out see if it fits? Thanks.

[westonpace]

This looks great, thanks for working on this. We will need to move decompression off the scheduling thread but I think we can do something simple for now and I have proposed an idea how we can handle it better in the long term.

[westonpace]

Yeah, this is a bit tricky, I see your concern. Unfortunately, we really shouldn't be doing the decompression here in the scheduler. The scheduler is single threaded and by doing compression here we will fall behind in scheduling I/O.

Something simple, that should work, is to wrap the compression task in a Arc<Mutex<...>> that each of the decode tasks that need that range share. Something roughly like...

struct ValuePageDecoder {
    bytes_per_value: u64,
    compressed_data: Vec<Bytes>,
    uncompressed_data: Arc<Mutex<Option<Bytes>>>,
}

impl ValuePageDecoder {

  // Called from decode_into
  fn get_uncompressed_bytes(&self) -> Result<Bytes> {
    let uncompressed_bytes = self.uncompressed_data.lock().unwrap();
    if uncompressed_bytes.is_none() {
      *uncompressed_bytes = self.decompress()?;
    }
    Ok(uncompressed_bytes.clone())
  }

}

This is not ideal because it means the decoder threads will be blocked for a bit but it should be good enough to start with.

A more complicated solution might be to have a dedicated thread pool for decompression. Once I/O is available we start to decompress it in this thread pool. I think we could do this as part of the I/O scheduler. So we can have arguments to submit_request which explain that the data is compressed and includes the parameters needed to decompress it. This would mean that compression is not just an "encoding" but something more fundamental but this is probably ok. We probably also want to treat encryption similarly. It also means we lose a bit of thread locality but I think that is inevitable. Decompression / decryption is costly enough that we want to parallelize across columns.

[westonpace]

We can use an environment variable on the write side. However, we should not use the environment variable on the read side. I think changes to encodings.proto are needed. For now, we can probably do this:

// Fixed width items placed contiguously in a buffer
message Flat {
  // the number of bits per value, must be greater than 0, does
  // not need to be a multiple of 8
  uint64 bits_per_value = 1;
  // the buffer of values
  Buffer buffer = 2;
  // The Compression message can specify the compression scheme (e.g. zstd) and any
  // other information that is needed for decompression.
  Compression compression = 3;
}

Later, if we end up going with the "dedicated thread pool for compression" and make compression part of the I/O scheduler then I think the Compression message will be a part of Buffer and not part of Flat. However, we can tackle that later.

[niyue]

I added such a structure in protobuf:

message Compression {
  string scheme = 1;
}

Currently, only the name of the compression scheme is stored and retrieved. Let me know if this is okay. Thanks.

[westonpace]

Strange. I thought we had something like this in the code base already but I cannot find it now.

[niyue]

@westonpace I updated the PR according to your review comments, could you please check it out? Thanks.

[westonpace]

Great, thanks for doing this. I've started moving towards "picking an encoder" in #2388 and it uses field metadata.

I've approved the tests if you want to try and look at any test failures, looks like there might be at least one format / lint issue.

Don't worry if you can't get to it by my tomorrow morning (~12 hours from now). I'll go ahead and fix anything and get this merged in at that time. I want to get this PR merged first so you don't have to deal with a rebase from #2388

[niyue]

I've updated the code, and hope the latest commit address all the lint issues. Thanks.

[niyue]

I have updated the code to address one remaining test issue.

I ran the tests locally, but one test case reported stepWithCompactEncoding - invalid compact unwind encoding, failed and stopped the test suite from completing. I believe this failure is unrelated to my changes. Please let me know if you think otherwise.

Thanks.

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 87.10938% with 33 lines in your changes are missing coverage. Please review.

Project coverage is 79.99%. Comparing base (8c1ee00) to head (d1403c3).
Report is 1 commits behind head on main.

Files	Patch %	Lines
...t/lance-encoding/src/encodings/physical/buffers.rs	73.68%	6 Missing and 4 partials ⚠️
rust/lance-encoding/src/encoder.rs	0.00%	6 Missing ⚠️
...ust/lance-encoding/src/encodings/physical/value.rs	96.52%	1 Missing and 3 partials ⚠️
rust/lance-file/src/v2/reader.rs	85.18%	2 Missing and 2 partials ⚠️
rust/lance-encoding/src/encodings/physical.rs	85.00%	3 Missing ⚠️
...ust/lance-encoding/src/encodings/physical/basic.rs	72.72%	0 Missing and 3 partials ⚠️
.../lance-encoding/src/encodings/logical/primitive.rs	71.42%	0 Missing and 2 partials ⚠️
...encoding/src/encodings/physical/fixed_size_list.rs	87.50%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2368      +/-   ##
==========================================
- Coverage   79.99%   79.99%   -0.01%     
==========================================
  Files         200      200              
  Lines       54519    54713     +194     
  Branches    54519    54713     +194     
==========================================
+ Hits        43612    43766     +154     
- Misses       8389     8410      +21     
- Partials     2518     2537      +19     

Flag	Coverage Δ
unittests	79.99% <87.10%> (-0.01%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

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