Skip to content

Technical Notes: Why Optimization Only Works with UncompressedChunk

Background

The three optimization methods (get_samples_continuous, get_samples_full, get_samples_batch_random_access) use byte-level access to directly read data from chunks. This approach provides significant performance improvements but has specific requirements.

Chunk Types in MULLER

1. UncompressedChunk

  • Data Storage: Raw, uncompressed bytes stored sequentially
  • Access Method: Direct byte offset calculation
  • data_bytes Property: Returns raw bytes directly from memory
  • Example: 
    Sample 0: bytes[0:100]
Sample 1: bytes[100:200]
Sample 2: bytes[200:300]

2. ChunkCompressedChunk

  • Data Storage: Entire chunk is compressed as a single unit
  • Access Method: Must decompress entire chunk first
  • data_bytes Property: Returns compressed bytes
  • Decompression: Uses decompressed_bytes or decompressed_samples after initialization
  • Why Optimization Fails:
  • chunk.data_bytes[start:end] returns compressed bytes, not raw data
  • Cannot use np.frombuffer() on compressed data
  • Must decompress entire chunk before accessing individual samples

3. SampleCompressedChunk

  • Data Storage: Each sample is independently compressed
  • Access Method: Must decompress each sample individually
  • data_bytes Property: Returns concatenated compressed samples
  • Why Optimization Fails:
  • Byte positions don't correspond to sample boundaries
  • Each sample has variable compressed size
  • Cannot calculate byte offset without decompressing

How Optimization Methods Work

_get_chunk_numpy_continuous()

def _get_chunk_numpy_continuous(chunk_engine, chunk_name, start_idx, end_idx):
    sample_dtype = chunk_engine.tensor_meta.dtype
    sample_size = np.array([], dtype=sample_dtype).itemsize
    chunk_key = get_chunk_key(chunk_engine.key, chunk_name)
    chunk = chunk_engine.get_chunk(chunk_key)

    # Direct byte-level access - ONLY works with uncompressed data
    np_bytes = chunk.data_bytes[start_idx*sample_size: end_idx*sample_size + sample_size]
    final_results = np.frombuffer(np_bytes, dtype=sample_dtype)
    return final_results

Key Assumptions: 1. chunk.data_bytes returns uncompressed raw bytes 2. Sample size is fixed and known (sample_size = dtype.itemsize) 3. Byte offset can be calculated as index * sample_size 4. np.frombuffer() can directly interpret bytes as numpy array

_get_chunk_numpy_full()

def _get_chunk_numpy_full(chunk_engine, chunk_name):
    sample_dtype = chunk_engine.tensor_meta.dtype
    chunk_key = get_chunk_key(chunk_engine.key, chunk_name)
    chunk = chunk_engine.get_chunk(chunk_key)

    # Direct access to all bytes - ONLY works with uncompressed data
    np_bytes = chunk.data_bytes
    final_results = np.frombuffer(np_bytes, dtype=sample_dtype)
    return final_results

Key Assumptions: 1. chunk.data_bytes contains all samples in uncompressed form 2. All samples are stored sequentially without gaps 3. Can directly convert entire byte array to numpy array

Why the Check is Necessary

Without the Check

If we allow compressed chunks to use these optimization methods:

# ChunkCompressedChunk example
chunk = ChunkCompressedChunk(...)
np_bytes = chunk.data_bytes[0:100]  # Returns compressed bytes!
result = np.frombuffer(np_bytes, dtype='float32')  # WRONG: interpreting compressed data as floats

Result: Garbage data, incorrect values, potential crashes

With the Check

def _validate_batch_samples(chunk_engine, fetch_chunks):
    return (
        fetch_chunks
        and not chunk_engine.is_video
        and not isinstance(chunk_engine.base_storage, MemoryProvider)
        and chunk_engine.chunk_compression is None  # Critical check
        and chunk_engine.sample_compression is None  # Critical check
    )

Result: Compressed chunks automatically fall back to get_samples(), which properly handles decompression

Performance Implications

UncompressedChunk (Optimization Enabled)

  • Full Access: 6.5x faster
  • Continuous Slice: 5x faster
  • Batch Random: 4x faster

Compressed Chunks (Fallback to get_samples())

  • No performance degradation compared to original implementation
  • Proper decompression ensures data correctness
  • Still benefits from parallel chunk loading when applicable

Implementation Details

Validation Location

chunk_engine_to_numpy_interface.py:214-228

Properties Used

  • chunk_engine.chunk_compression: Returns compression algorithm if chunk-wise compressed, else None
  • chunk_engine.sample_compression: Returns compression algorithm if sample-wise compressed, else None

Fallback Behavior

When validation fails, the system uses get_samples() which: 1. Properly handles all chunk types 2. Decompresses data as needed 3. Maintains data integrity 4. Provides correct results for all compression types

Conclusion

The restriction to UncompressedChunk is not a limitation but a correctness requirement. The optimization methods rely on direct byte-level access, which is only valid for uncompressed data. The validation check ensures:

  1. Correctness: Compressed data is never misinterpreted as raw data
  2. Performance: Uncompressed chunks get maximum optimization
  3. Compatibility: Compressed chunks still work correctly via fallback
  4. Safety: No risk of data corruption or incorrect results