mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-23 20:22:37 +00:00
lancedb-gatefixer[bot]
624defd109
## Root cause The pinned Lance Arrow crate rebuilt a sliced struct child with its original Arrow offset but supplied a validity bitmap containing only the visible rows. Arrow requires the bitmap to cover the offset as well, so filtered vector searches involving nullable nested fields could panic with `null_bit_buffer size too small`. ## Fix - Carry an offset-aligned validity bitmap when propagating parent nulls to a sliced child. - Vendor the pinned `lance-arrow` crate until a Lance release includes the fix; include it in Maturin source distributions. - Add unit regressions for sliced children and parent bitmaps, plus a record-batch merge regression. ## Validation - The issue's 8,600-row Python reproducer completes its filtered flat dot search and returns 8,448 rows without a panic. - The targeted Rust merge regression and two focused Python query tests pass. - Rust formatting, Ruff formatting and lint, the vendored crate tests, and the workspace Cargo checks passed. The Maturin source archive includes the patched crate. Fixes #4260 <!-- lance-gatekeeper-fix:v1 agent=e6045f701d9b708ee26528ba36353aba generation=1 --> --------- Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com>
93 lines
2.7 KiB
Rust
93 lines
2.7 KiB
Rust
// SPDX-License-Identifier: Apache-2.0
|
|
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
|
// SPDX-FileCopyrightText: Copyright The Lance Authors
|
|
|
|
use std::collections::HashSet;
|
|
|
|
use arrow_array::{Array, RecordBatch};
|
|
use arrow_data::ArrayData;
|
|
|
|
/// Counts memory used by buffers of Arrow arrays and RecordBatches.
|
|
///
|
|
/// This is meant to capture how much memory is being used by the Arrow data
|
|
/// structures as they are. It does not represent the memory used if the data
|
|
/// were to be serialized and then deserialized. In particular:
|
|
///
|
|
/// * This does not double count memory used by buffers shared by multiple
|
|
/// arrays or batches. Round-tripped data may use more memory because of this.
|
|
/// * This counts the **total** size of the buffers, even if the array is a slice.
|
|
/// Round-tripped data may use less memory because of this.
|
|
#[derive(Default)]
|
|
pub struct MemoryAccumulator {
|
|
seen: HashSet<usize>,
|
|
total: usize,
|
|
}
|
|
|
|
impl MemoryAccumulator {
|
|
pub fn record_array(&mut self, array: &dyn Array) {
|
|
let data = array.to_data();
|
|
self.record_array_data(&data);
|
|
}
|
|
|
|
fn record_array_data(&mut self, data: &ArrayData) {
|
|
for buffer in data.buffers() {
|
|
let ptr = buffer.as_ptr();
|
|
if self.seen.insert(ptr as usize) {
|
|
self.total += buffer.capacity();
|
|
}
|
|
}
|
|
|
|
if let Some(nulls) = data.nulls() {
|
|
let null_buf = nulls.inner().inner();
|
|
let ptr = null_buf.as_ptr();
|
|
if self.seen.insert(ptr as usize) {
|
|
self.total += null_buf.capacity();
|
|
}
|
|
}
|
|
|
|
for child in data.child_data() {
|
|
self.record_array_data(child);
|
|
}
|
|
}
|
|
|
|
pub fn record_batch(&mut self, batch: &RecordBatch) {
|
|
for array in batch.columns() {
|
|
self.record_array(array);
|
|
}
|
|
}
|
|
|
|
pub fn total(&self) -> usize {
|
|
self.total
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use std::sync::Arc;
|
|
|
|
use arrow_array::Int32Array;
|
|
use arrow_schema::{DataType, Field, Schema};
|
|
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_memory_accumulator() {
|
|
let batch = RecordBatch::try_new(
|
|
Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
|
|
vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
|
|
)
|
|
.unwrap();
|
|
let slice = batch.slice(1, 2);
|
|
|
|
let mut acc = MemoryAccumulator::default();
|
|
|
|
// Should record whole buffer, not just slice
|
|
acc.record_batch(&slice);
|
|
assert_eq!(acc.total(), 3 * std::mem::size_of::<i32>());
|
|
|
|
// Should not double count
|
|
acc.record_batch(&slice);
|
|
assert_eq!(acc.total(), 3 * std::mem::size_of::<i32>());
|
|
}
|
|
}
|