Repository: pola-rs/polars
File: crates/polars-arrow/src/io/ipc/read/read_basic.rs (Rust)

Code before:
    scratch.clear();
    scratch.try_reserve(buffer_length)?;
    reader
        .by_ref()
        .take(buffer_length as u64)
        .read_to_end(scratch)?;

    polars_ensure!(scratch.len() == buffer_length, ComputeError: "Malformed IPC file: expected compressed buffer of len {buffer_length}, got {}", scratch.len());

    let decompressed_len_field = i64::from_le_bytes(scratch[..8].try_into().unwrap());
    let decompressed_bytes: usize = if decompressed_len_field == -1 {
        buffer_length - 8
    } else {
        decompressed_len_field.try_into().map_err(|_| {
            polars_err!(ComputeError: "Malformed IPC file: got invalid decompressed length {decompressed_len_field}")
        })?
    };

    polars_ensure!(decompressed_bytes.is_multiple_of(size_of::<T>()),
            ComputeError: "Malformed IPC file: got decompressed buffer length which is not a multiple of the data type");
    let total_length_in_array = decompressed_bytes / size_of::<T>();

    // Note: in case of slicing, the lengths may not match
    if let Some(row_limit) = row_limit {
        polars_ensure!(row_limit <= total_length_in_array, 
            ComputeError: "Malformed IPC file: got unexpected decompressed buffer size {total_length_in_array}, expected {row_limit}");
    }

    if decompressed_len_field == -1 {
        return Ok(bytemuck::cast_slice(&scratch[8..]).to_vec());
    }

    // It is undefined behavior to call read_exact on un-initialized, https://doc.rust-lang.org/std/io/trait.Read.html#tymethod.read
    // see also https://github.com/MaikKlein/ash/issues/354#issue-781730580
    let exact_row_size = match row_limit {
        Some(row_limit) => row_limit,
        None => total_length_in_array,
    };

    let mut buffer = vec![T::default(); exact_row_size];
    let out_slice = bytemuck::cast_slice_mut(&mut buffer);

    let compression = compression
        .codec()
        .map_err(|err| polars_err!(oos = OutOfSpecKind::InvalidFlatbufferCompression(err)))?;

    match compression {
        arrow_format::ipc::CompressionType::Lz4Frame => {
            compression::decompress_lz4(&scratch[8..], out_slice)?;
        },
        arrow_format::ipc::CompressionType::Zstd => {

Review comment (suggestion):
This technically doesn't imply a malformed IPC file - I think the error should instead say "too many rows requested".

We should also rename `row_limit` to `n_rows` - it's technically not a limit but rather an exact request.


Code after:
    // Decompress first.
    scratch.clear();
    scratch.try_reserve(buffer_length)?;
    reader
        .by_ref()
        .take(buffer_length as u64)
        .read_to_end(scratch)?;

    polars_ensure!(scratch.len() == buffer_length, ComputeError: "Malformed IPC file: expected compressed buffer of len {buffer_length}, got {}", scratch.len());

    let decompressed_len_field = i64::from_le_bytes(scratch[..8].try_into().unwrap());
    let decompressed_bytes: usize = if decompressed_len_field == -1 {
        buffer_length - 8
    } else {
        decompressed_len_field.try_into().map_err(|_| {
            polars_err!(ComputeError: "Malformed IPC file: got invalid decompressed length {decompressed_len_field}")
        })?
    };

    polars_ensure!(decompressed_bytes.is_multiple_of(size_of::<T>()),
            ComputeError: "Malformed IPC file: got decompressed buffer length which is not a multiple of the data type");
    let n_rows_in_array = decompressed_bytes / size_of::<T>();

    if decompressed_len_field == -1 {
        return Ok(bytemuck::cast_slice(&scratch[8..]).to_vec());
    }

    // It is undefined behavior to call read_exact on un-initialized, https://doc.rust-lang.org/std/io/trait.Read.html#tymethod.read
    // see also https://github.com/MaikKlein/ash/issues/354#issue-781730580

    let n_rows_exact = row_limit
        .map(|limit| std::cmp::min(limit, n_rows_in_array))
        .unwrap_or(n_rows_in_array);

    let mut buffer = vec![T::default(); n_rows_exact];
    let out_slice = bytemuck::cast_slice_mut(&mut buffer);

    let compression = compression
        .codec()
        .map_err(|err| polars_err!(oos = OutOfSpecKind::InvalidFlatbufferCompression(err)))?;

    match compression {
        arrow_format::ipc::CompressionType::Lz4Frame => {
            compression::decompress_lz4(&scratch[8..], out_slice)?;
        },
        arrow_format::ipc::CompressionType::Zstd => {
            compression::decompress_zstd(&scratch[8..], out_slice)?;
        },
    }
    Ok(buffer)
}