diff --git a/Cargo.lock b/Cargo.lock index 6691c1efd5c..2b7b658074f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11451,11 +11451,13 @@ dependencies = [ "mimalloc", "num-traits", "rstest", + "rustc-hash", "simdutf8", "tracing", "vortex-array", "vortex-buffer", "vortex-error", + "vortex-fastlanes", "vortex-fsst", "vortex-mask", "vortex-onpair", diff --git a/encodings/fastlanes/src/bitpacking/compute/cast.rs b/encodings/fastlanes/src/bitpacking/compute/cast.rs index 8203e68a6b7..53f0e8d8481 100644 --- a/encodings/fastlanes/src/bitpacking/compute/cast.rs +++ b/encodings/fastlanes/src/bitpacking/compute/cast.rs @@ -32,6 +32,14 @@ fn is_widening_int_cast(src: PType, tgt: PType) -> bool { && (src.is_unsigned_int() || tgt.is_signed_int()) } +/// Returns `true` if every value a patch-free array packed at `bit_width` can hold fits `tgt`. +/// Packed values are non-negative and below `2^bit_width`. +fn packed_values_fit(bit_width: u8, src: PType, tgt: PType) -> bool { + src.is_int() + && tgt.is_int() + && usize::from(bit_width) + usize::from(tgt.is_signed_int()) <= tgt.bit_width() +} + fn build_with_validity( array: ArrayView<'_, BitPacked>, bit_width: u8, @@ -89,15 +97,17 @@ impl CastKernel for BitPacked { return build_with_validity(array, bit_width, dtype, new_validity).map(Some); } - // Widening integer cast: unpack each FastLanes chunk into a cache-resident scratch buffer - // and cast-copy straight into the wide output, avoiding a full-length intermediate buffer - // and the generic cast kernel's bounds-check scan (unnecessary when widening). + // Widening integer cast, or a narrowing one whose target holds every packed value: unpack + // each FastLanes chunk into a cache-resident scratch buffer and cast-copy straight into the + // output, avoiding a full-length intermediate buffer and the generic cast kernel's + // bounds-check scan (unnecessary when every value fits). let DType::Primitive(tgt, tgt_nullability) = dtype else { return Ok(None); }; let (tgt, tgt_nullability) = (*tgt, *tgt_nullability); let src = array.dtype().as_ptype(); - if !is_widening_int_cast(src, tgt) { + let fits = || array.patches().is_none() && packed_values_fit(bit_width, src, tgt); + if !is_widening_int_cast(src, tgt) && !fits() { return Ok(None); } @@ -144,6 +154,7 @@ mod tests { use vortex_session::VortexSession; use crate::BitPackedArray; + use crate::BitPackedArrayExt; use crate::BitPackedData; static SESSION: LazyLock = LazyLock::new(|| { @@ -266,6 +277,39 @@ mod tests { Ok(()) } + /// Narrowing casts of patch-free arrays whose bit width fits the target unpack directly and + /// match a plain primitive cast, including sliced and nullable inputs. + #[rstest] + #[case(PType::I64, PType::U8)] + #[case(PType::I64, PType::I8)] + #[case(PType::U64, PType::U16)] + #[case(PType::I32, PType::I16)] + fn test_cast_bitpacked_narrowing_fits( + #[case] src: PType, + #[case] tgt: PType, + ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let source = match_each_integer_ptype!(src, |S| { + PrimitiveArray::from_option_iter((0..2051usize).map(|i| { + (i % 11 != 5).then(|| { + ::from_usize(i % 128) + .expect("test values fit every integer ptype") + }) + })) + }) + .into_array(); + let packed = bp(&source, 7); + assert!(packed.patches().is_none()); + let target = DType::Primitive(tgt, Nullability::Nullable); + + let casted = packed.clone().into_array().cast(target.clone())?; + assert_arrays_eq!(casted, source.cast(target.clone())?, &mut ctx); + + let sliced = packed.into_array().slice(500..1700)?.cast(target.clone())?; + assert_arrays_eq!(sliced, source.slice(500..1700)?.cast(target)?, &mut ctx); + Ok(()) + } + #[rstest] #[case(bp(&buffer![0u8, 10, 20, 30, 40, 50, 60, 63].into_array(), 6))] #[case(bp(&buffer![0u16, 100, 200, 300, 400, 500].into_array(), 9))] diff --git a/vortex-arrow/Cargo.toml b/vortex-arrow/Cargo.toml index 8605b37a3f0..948e03a9865 100644 --- a/vortex-arrow/Cargo.toml +++ b/vortex-arrow/Cargo.toml @@ -27,11 +27,13 @@ arrow-schema = { workspace = true, features = ["canonical_extension_types"] } arrow-select = { workspace = true } itertools = { workspace = true } num-traits = { workspace = true } +rustc-hash = { workspace = true } simdutf8 = { workspace = true } tracing = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true, features = ["arrow"] } vortex-error = { workspace = true } +vortex-fastlanes = { workspace = true } vortex-mask = { workspace = true } vortex-runend = { workspace = true } vortex-session = { workspace = true } diff --git a/vortex-arrow/src/executor/dictionary.rs b/vortex-arrow/src/executor/dictionary.rs index 343c08357af..d0e7c96352d 100644 --- a/vortex-arrow/src/executor/dictionary.rs +++ b/vortex-arrow/src/executor/dictionary.rs @@ -1,15 +1,24 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::collections::hash_map::Entry; +use std::hash::Hash; use std::sync::Arc; +use arrow_array::Array; use arrow_array::ArrayRef as ArrowArrayRef; +use arrow_array::ArrowPrimitiveType; use arrow_array::DictionaryArray; use arrow_array::PrimitiveArray; use arrow_array::cast::AsArray; use arrow_array::new_null_array; use arrow_array::types::*; +use arrow_buffer::ArrowNativeType; +use arrow_buffer::ScalarBuffer; use arrow_schema::DataType; +use arrow_schema::TimeUnit; +use rustc_hash::FxHashMap; +use rustc_hash::FxHashSet; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; @@ -17,11 +26,19 @@ use vortex_array::arrays::Constant; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::Dict; use vortex_array::arrays::DictArray; +use vortex_array::arrays::PrimitiveArray as VortexPrimitiveArray; use vortex_array::arrays::dict::DictArraySlotsExt; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::PType; +use vortex_array::match_each_integer_ptype; use vortex_array::matcher::Matcher; use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_fastlanes::BitPacked; +use vortex_fastlanes::BitPackedArrayExt; +use vortex_fastlanes::BitWidthsView; use crate::ArrowExporter; use crate::executor::execute_arrow_naive; @@ -33,7 +50,7 @@ impl Matcher for ArrowDictExportable { type Match<'a> = &'a ArrayRef; fn try_match(array: &ArrayRef) -> Option> { - (array.is::() || array.is::()).then_some(array) + (array.is::() || array.is::() || array.is::()).then_some(array) } } @@ -54,10 +71,16 @@ pub(super) fn to_arrow_dictionary( Ok(constant) => return constant_to_dict(constant, codes_type, values_type, exporter, ctx), Err(array) => array, }; + if let Some(dict) = bitpacked_to_dict(&array, codes_type, values_type, exporter, ctx)? { + return Ok(dict); + } // Otherwise, we should try and build a dictionary. - // Arrow hides this functionality inside the cast module! let array = execute_arrow_naive(array, Some(values_type), exporter, ctx)?; + if let Some(dict) = primitive_to_dict(&array, codes_type) { + return Ok(dict); + } + // Arrow hides this functionality inside the cast module! arrow_cast::cast( &array, &DataType::Dictionary(Box::new(codes_type.clone()), Box::new(values_type.clone())), @@ -65,6 +88,274 @@ pub(super) fn to_arrow_dictionary( .map_err(VortexError::from) } +/// The widest packed values [`bitpacked_to_dict`] exports as codes. +const MAX_BITPACKED_CODE_BITS: u8 = 16; + +/// Exports a patch-free bit-packed array as a dictionary over every value its bit width can hold. +/// +/// Packed values are non-negative and below `2^bit_width`, so the unpacked values are their own +/// codes into the values `0..2^bit_width`. Unpacking straight into the codes type skips decoding +/// to the values type and then dictionary-encoding the result. Returns `None` when the array is not +/// such a bit-packed array, the codes type cannot hold every packed value, or the dictionary would +/// outnumber the array's rows. +fn bitpacked_to_dict( + array: &ArrayRef, + codes_type: &DataType, + values_type: &DataType, + exporter: &ArrowExporter<'_>, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let Some(bitpacked) = array.as_opt::() else { + return Ok(None); + }; + let BitWidthsView::Global(bit_width) = bitpacked.bit_widths() else { + return Ok(None); + }; + let Some(codes_ptype) = codes_ptype(codes_type) else { + return Ok(None); + }; + let dict_len = 1usize << bit_width; + if bit_width > MAX_BITPACKED_CODE_BITS + || usize::from(bit_width) + usize::from(codes_ptype.is_signed_int()) + > codes_ptype.bit_width() + || !values_type.is_integer() + || dict_len > array.len() + || bitpacked.patches().is_some() + { + return Ok(None); + } + + let codes = array + .clone() + .cast(DType::Primitive(codes_ptype, array.dtype().nullability()))?; + let codes = execute_arrow_naive(codes, Some(codes_type), exporter, ctx)?; + let values = match_each_integer_ptype!(array.dtype().as_ptype(), |T| { + VortexPrimitiveArray::from_iter((0..dict_len).map(T::usize_as)).into_array() + }); + let values = execute_arrow_naive(values, Some(values_type), exporter, ctx)?; + make_dict_array(codes_type, codes, values).map(Some) +} + +fn codes_ptype(codes_type: &DataType) -> Option { + Some(match codes_type { + DataType::Int8 => PType::I8, + DataType::Int16 => PType::I16, + DataType::Int32 => PType::I32, + DataType::Int64 => PType::I64, + DataType::UInt8 => PType::U8, + DataType::UInt16 => PType::U16, + DataType::UInt32 => PType::U32, + DataType::UInt64 => PType::U64, + _ => return None, + }) +} + +/// Dictionary-encodes a primitive Arrow array with codes of `codes_type`. +/// +/// Arrow's cast hashes every value with SipHash. This hashes the values' bits with FxHash instead, +/// which dominates exporting large nested columns that the reader requests as dictionaries. +/// Returns `None` for unsupported types and when the distinct values overflow the codes type, so +/// the caller can fall back to Arrow's cast. +fn primitive_to_dict(array: &ArrowArrayRef, codes_type: &DataType) -> Option { + match codes_type { + DataType::Int8 => primitive_to_dict_with::(array), + DataType::Int16 => primitive_to_dict_with::(array), + DataType::Int32 => primitive_to_dict_with::(array), + DataType::Int64 => primitive_to_dict_with::(array), + DataType::UInt8 => primitive_to_dict_with::(array), + DataType::UInt16 => primitive_to_dict_with::(array), + DataType::UInt32 => primitive_to_dict_with::(array), + DataType::UInt64 => primitive_to_dict_with::(array), + _ => None, + } +} + +fn primitive_to_dict_with( + array: &ArrowArrayRef, +) -> Option { + match array.data_type() { + DataType::Int8 => encode_primitive::(array, |v| v), + DataType::Int16 => encode_primitive::(array, |v| v), + DataType::Int32 => encode_primitive::(array, |v| v), + DataType::Int64 => encode_primitive::(array, |v| v), + DataType::UInt8 => encode_primitive::(array, |v| v), + DataType::UInt16 => encode_primitive::(array, |v| v), + DataType::UInt32 => encode_primitive::(array, |v| v), + DataType::UInt64 => encode_primitive::(array, |v| v), + DataType::Float16 => encode_primitive::(array, |v| v.to_bits()), + DataType::Float32 => encode_primitive::(array, |v| v.to_bits()), + DataType::Float64 => encode_primitive::(array, |v| v.to_bits()), + DataType::Date32 => encode_primitive::(array, |v| v), + DataType::Date64 => encode_primitive::(array, |v| v), + DataType::Timestamp(TimeUnit::Second, _) => { + encode_primitive::(array, |v| v) + } + DataType::Timestamp(TimeUnit::Millisecond, _) => { + encode_primitive::(array, |v| v) + } + DataType::Timestamp(TimeUnit::Microsecond, _) => { + encode_primitive::(array, |v| v) + } + DataType::Timestamp(TimeUnit::Nanosecond, _) => { + encode_primitive::(array, |v| v) + } + _ => None, + } +} + +/// Valid values sampled to choose how [`encode_primitive`] assigns codes. +const SAMPLE_LEN: usize = 1024; +/// The largest key range [`encode_primitive`] looks up in a dense table. +const MAX_DENSE_RANGE: usize = 1 << 16; + +/// Dictionary-encodes `array` keyed by `key` of each valid value. Null values get null codes. +/// +/// Keys spanning a small range are looked up in a dense table. When a sample of the values is +/// nearly all distinct, the values themselves become the dictionary with sequential codes, since +/// Arrow dictionaries may hold duplicates and deduplicating would barely shrink them. Otherwise +/// codes are assigned in order of first appearance through a hash map. +fn encode_primitive( + array: &ArrowArrayRef, + key: impl Fn(T::Native) -> B, +) -> Option +where + K: ArrowDictionaryKeyType, + T: ArrowPrimitiveType, + B: Hash + Eq + Copy + Into, +{ + let array = array.as_primitive::(); + let nulls = array.nulls().filter(|nulls| nulls.null_count() > 0); + let is_valid = |idx: usize| nulls.is_none_or(|nulls| nulls.is_valid(idx)); + let valid_keys = || { + array + .values() + .iter() + .enumerate() + .filter(move |(idx, _)| is_valid(*idx)) + .map(|(_, &value)| key(value)) + }; + + // A sample decides how to assign codes, and rules out a dense table without scanning every + // key when the sampled keys already span too wide a range. + let mut sample = FxHashSet::default(); + let (mut sample_min, mut sample_max) = (i128::MAX, i128::MIN); + let sampled = valid_keys() + .take(SAMPLE_LEN) + .inspect(|key| { + sample.insert(*key); + let key = (*key).into(); + sample_min = sample_min.min(key); + sample_max = sample_max.max(key); + }) + .count(); + let dense = sample_max + .checked_sub(sample_min) + .is_some_and(|range| range < MAX_DENSE_RANGE as i128) + .then(|| { + let (min, max) = valid_keys().fold((i128::MAX, i128::MIN), |(min, max), key| { + let key = key.into(); + (min.min(key), max.max(key)) + }); + let range = usize::try_from(max - min).ok()?; + (range < MAX_DENSE_RANGE && range <= array.len().saturating_mul(4)) + .then_some((min, range)) + }) + .flatten(); + + let (codes, values) = if let Some((min, range)) = dense { + encode_dense::(array, &is_valid, &key, min, range)? + } else if sampled > 0 + && sample.len() * 8 >= sampled * 7 + && K::Native::from_usize(array.len()).is_some() + { + let codes = (0..array.len()).map(K::Native::usize_as).collect(); + (codes, array.values().clone()) + } else { + let distinct = sample.len() * array.len() / sampled.max(1); + encode_hashed::(array, &is_valid, &key, distinct)? + }; + + let codes = PrimitiveArray::::new(codes, nulls.cloned()); + let values = PrimitiveArray::::new(values, None).with_data_type(array.data_type().clone()); + // SAFETY: every valid code indexes a value of the dictionary. + Some(Arc::new(unsafe { + DictionaryArray::new_unchecked(codes, Arc::new(values)) + })) +} + +/// Codes and dictionary values built by [`encode_dense`] and [`encode_hashed`]. +type EncodedDict = ( + ScalarBuffer<::Native>, + ScalarBuffer<::Native>, +); + +/// Assigns codes in order of first appearance through a table indexed by `key - min`. +fn encode_dense( + array: &PrimitiveArray, + is_valid: &impl Fn(usize) -> bool, + key: &impl Fn(T::Native) -> B, + min: i128, + range: usize, +) -> Option> +where + K: ArrowDictionaryKeyType, + T: ArrowPrimitiveType, + B: Into, +{ + const UNSEEN: u32 = u32::MAX; + let mut table = vec![UNSEEN; range + 1]; + let mut values = Vec::new(); + let mut codes = Vec::with_capacity(array.len()); + for (idx, &value) in array.values().iter().enumerate() { + if !is_valid(idx) { + codes.push(K::Native::default()); + continue; + } + let slot = &mut table[usize::try_from(key(value).into() - min).ok()?]; + if *slot == UNSEEN { + K::Native::from_usize(values.len())?; + *slot = u32::try_from(values.len()).ok()?; + values.push(value); + } + codes.push(K::Native::usize_as(*slot as usize)); + } + Some((codes.into(), values.into())) +} + +/// Assigns codes in order of first appearance through a hash map sized for `distinct` keys. +fn encode_hashed( + array: &PrimitiveArray, + is_valid: &impl Fn(usize) -> bool, + key: &impl Fn(T::Native) -> B, + distinct: usize, +) -> Option> +where + K: ArrowDictionaryKeyType, + T: ArrowPrimitiveType, + B: Hash + Eq, +{ + let mut lookup: FxHashMap = + FxHashMap::with_capacity_and_hasher(distinct, Default::default()); + let mut values = Vec::with_capacity(distinct); + let mut codes = Vec::with_capacity(array.len()); + for (idx, &value) in array.values().iter().enumerate() { + if !is_valid(idx) { + codes.push(K::Native::default()); + continue; + } + let code = match lookup.entry(key(value)) { + Entry::Occupied(entry) => *entry.get(), + Entry::Vacant(entry) => { + let code = K::Native::from_usize(values.len())?; + values.push(value); + *entry.insert(code) + } + }; + codes.push(code); + } + Some((codes.into(), values.into())) +} + /// Convert a constant array to a dictionary with a single entry. fn constant_to_dict( array: ConstantArray, @@ -160,7 +451,14 @@ mod tests { use std::sync::LazyLock; use arrow_array::DictionaryArray as ArrowDictArray; + use arrow_array::Float64Array; + use arrow_array::Int16Array; + use arrow_array::Int64Array; use arrow_array::StringArray; + use arrow_array::UInt32Array; + use arrow_array::UInt64Array; + use arrow_array::cast::AsArray; + use arrow_array::types::Int16Type; use arrow_array::types::UInt8Type; use arrow_array::types::UInt32Type; use arrow_schema::DataType; @@ -176,6 +474,8 @@ mod tests { use vortex_array::scalar_fn::fns::mask::Mask; use vortex_buffer::buffer; use vortex_error::VortexResult; + use vortex_fastlanes::BitPackedArrayExt; + use vortex_fastlanes::BitPackedData; use vortex_session::VortexSession; use crate::ArrowArrayExecutor; @@ -239,6 +539,22 @@ mod tests { dict_type(DataType::UInt8, DataType::Utf8), Arc::new(vec![Some("a"), None, Some("a"), Some("b"), Some("a")].into_iter().collect::>()) as arrow_array::ArrayRef, )] + #[case::primitive_fallback( + PrimitiveArray::from_option_iter([Some(7i64), None, Some(-3), Some(7), Some(-3)]).into_array(), + dict_type(DataType::UInt32, DataType::Int64), + Arc::new(ArrowDictArray::::new( + UInt32Array::from(vec![Some(0), None, Some(1), Some(0), Some(1)]), + Arc::new(Int64Array::from(vec![7, -3])), + )) as arrow_array::ArrayRef, + )] + #[case::float_fallback_by_bits( + buffer![0.0f64, -0.0, f64::NAN, 0.0, f64::NAN].into_array(), + dict_type(DataType::Int16, DataType::Float64), + Arc::new(ArrowDictArray::::new( + Int16Array::from(vec![0, 1, 2, 0, 2]), + Arc::new(Float64Array::from(vec![0.0, -0.0, f64::NAN])), + )) as arrow_array::ArrayRef, + )] fn to_arrow_dictionary( #[case] input: vortex_array::ArrayRef, #[case] target_type: DataType, @@ -249,6 +565,59 @@ mod tests { Ok(()) } + #[test] + fn primitive_fallback_distinct_values() -> VortexResult<()> { + let values = + (0..2000u64).map(|i| (i % 7 != 3).then(|| i.wrapping_mul(0x9E37_79B9_7F4A_7C15))); + let input = PrimitiveArray::from_option_iter(values.clone()).into_array(); + let actual = execute(input, &dict_type(DataType::UInt32, DataType::UInt64))?; + let expected = UInt64Array::from_iter(values); + let flat = arrow_cast::cast(&actual, &DataType::UInt64)?; + assert_eq!(flat.as_ref(), &expected as &dyn arrow_array::Array); + Ok(()) + } + + #[rstest] + #[case::fits(DataType::UInt32, 11)] + #[case::signed_codes(DataType::Int16, 11)] + #[case::codes_narrower_than_bit_width(DataType::UInt16, 20)] + #[case::wide(DataType::UInt32, 20)] + fn bitpacked_exports_as_dict( + #[case] codes_type: DataType, + #[case] bit_width: u8, + ) -> VortexResult<()> { + let max = (1u64 << bit_width) - 1; + let values: Vec<_> = (0..5000u64) + .map(|i| (i % 13 != 4).then(|| i64::try_from(i * 7919 % max).unwrap_or(0))) + .collect(); + let source = PrimitiveArray::from_option_iter(values.clone()).into_array(); + let packed = + BitPackedData::encode(&source, bit_width, &mut SESSION.create_execution_ctx())?; + assert!(packed.patches().is_none()); + + let actual = execute(packed.into_array(), &dict_type(codes_type, DataType::Int64))?; + let flat = arrow_cast::cast(&actual, &DataType::Int64)?; + assert_eq!( + flat.as_ref(), + &Int64Array::from(values) as &dyn arrow_array::Array + ); + Ok(()) + } + + #[test] + fn primitive_fallback_fills_codes_type() -> VortexResult<()> { + let fits = PrimitiveArray::from_iter(0..=u8::MAX).into_array(); + let actual = execute(fits, &dict_type(DataType::UInt8, DataType::UInt8))?; + assert_eq!(actual.as_dictionary::().values().len(), 256); + Ok(()) + } + + #[test] + fn primitive_fallback_overflowing_codes_type_errors() { + let overflows = PrimitiveArray::from_iter(0..=256u16).into_array(); + assert!(execute(overflows, &dict_type(DataType::UInt8, DataType::UInt16)).is_err()); + } + #[test] fn mask_wrapped_dict_exports() -> VortexResult<()> { // Dictionary behind a lazy `mask` scalar-fn — the shape a scan produces when a row diff --git a/vortex-layout/src/scan/execution/projection.rs b/vortex-layout/src/scan/execution/projection.rs index 00c0a0614f6..71bf4c5c653 100644 --- a/vortex-layout/src/scan/execution/projection.rs +++ b/vortex-layout/src/scan/execution/projection.rs @@ -1,8 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::collections::VecDeque; use std::vec::IntoIter; +use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::arrays::ChunkedArray; use vortex_array::dtype::DType; @@ -30,6 +32,8 @@ use crate::scan::execution::adaptor::Piece; use crate::scan::execution::adaptor::ScanAdaptor; const MAX_OUTPUT_ROWS: usize = 65_536; +/// Pieces smaller than this are grouped with their neighbours into one batch. +const MIN_OUTPUT_ROWS: usize = 4_096; /// Keeps sparse selections together and subdivides dense selections to avoid large copies when /// converting chunked projection results to Arrow. @@ -97,15 +101,19 @@ impl Planner for ProjectionPlanner { } } -/// Runs the projection plan over a split's selected rows and returns them as one batch, in row -/// order. A selection that projects to no rows finishes without a batch. +/// Runs the projection plan over a split's selected rows and returns them in row order, one +/// batch per piece the scan produced. Pieces under [`MIN_OUTPUT_ROWS`] are grouped with their +/// neighbours. Emitting pieces separately avoids concatenating them on conversion to Arrow, +/// which rebuilds nested dictionaries. A selection that projects to no rows finishes without a +/// batch. struct ProjectionMorsel { plans: Option, dtype: DType, selected: SelectedRows, adaptor: Option, pieces: Vec, - done: bool, + /// Batches ready to emit once the scan has finished, in row order. + batches: Option>, } impl ProjectionMorsel { @@ -117,24 +125,39 @@ impl ProjectionMorsel { selected, adaptor: None, pieces: Vec::new(), - done: false, + batches: None, } } - fn finish(&mut self) -> VortexResult { - self.done = true; + /// Groups the finished pieces into batches of at least [`MIN_OUTPUT_ROWS`] rows where + /// possible. + fn group_pieces(&mut self) -> VortexResult> { self.pieces.sort_by_key(|piece| piece.rows.start); - let mut arrays: Vec<_> = std::mem::take(&mut self.pieces) - .into_iter() - .map(|piece| piece.array) - .filter(|array| !array.is_empty()) - .collect(); - let array = match arrays.len() { - 0 => return Ok(MorselOutput::Done), - 1 => arrays.remove(0), - _ => ChunkedArray::try_new(arrays, self.dtype.clone())?.into_array(), - }; - Ok(MorselOutput::Batch(array)) + let mut batches = VecDeque::new(); + let mut group = Vec::new(); + let mut group_rows = 0; + for piece in std::mem::take(&mut self.pieces) { + if piece.array.is_empty() { + continue; + } + group_rows += piece.array.len(); + group.push(piece.array); + if group_rows >= MIN_OUTPUT_ROWS { + batches.push_back(self.concat(std::mem::take(&mut group))?); + group_rows = 0; + } + } + if !group.is_empty() { + batches.push_back(self.concat(group)?); + } + Ok(batches) + } + + fn concat(&self, mut arrays: Vec) -> VortexResult { + if arrays.len() == 1 { + return Ok(arrays.remove(0)); + } + Ok(ChunkedArray::try_new(arrays, self.dtype.clone())?.into_array()) } } @@ -148,8 +171,14 @@ impl IoConsumer for ProjectionMorsel { impl Morsel for ProjectionMorsel { fn state(&self) -> State { + if let Some(batches) = &self.batches { + return if batches.is_empty() { + State::Done + } else { + State::NeedsCompute + }; + } match &self.adaptor { - _ if self.done => State::Done, None => State::NeedsCompute, Some(adaptor) => match adaptor.state() { State::Done => State::NeedsCompute, @@ -159,8 +188,11 @@ impl Morsel for ProjectionMorsel { } fn compute(&mut self) -> VortexResult { - if self.done { - vortex_bail!("ProjectionMorsel: compute called after Done"); + if let Some(batches) = &mut self.batches { + return match batches.pop_front() { + Some(batch) => Ok(MorselOutput::Batch(batch)), + None => vortex_bail!("ProjectionMorsel: compute called after Done"), + }; } let Some(adaptor) = self.adaptor.as_mut() else { let plans = self @@ -185,7 +217,8 @@ impl Morsel for ProjectionMorsel { return Ok(MorselOutput::Continue); }; if adaptor.state() == State::Done { - return self.finish(); + self.batches = Some(self.group_pieces()?); + return Ok(MorselOutput::Continue); } Ok(match adaptor.compute()? { AdaptorStep::Yield => MorselOutput::Continue, diff --git a/vortex-layout/src/scan/v2/cache.rs b/vortex-layout/src/scan/v2/cache.rs index 3daa0975a0e..5a02816a00d 100644 --- a/vortex-layout/src/scan/v2/cache.rs +++ b/vortex-layout/src/scan/v2/cache.rs @@ -9,6 +9,7 @@ use parking_lot::Mutex; use rustc_hash::FxHashMap; use vortex_array::expr::BoundExpression; use vortex_array::expr::ExactBoundExpr; +use vortex_error::VortexResult; use super::io::SegmentRanges; use crate::plan::physical::PlanRef; @@ -36,6 +37,8 @@ struct Inner { /// them, so a projection must reuse its plan for its values to be found again. projections: Mutex>, ranges: Mutex>, + /// Serialized bytes the whole-file projection reads, by projection. + projection_bytes: Mutex>, values: ValueCache, } @@ -80,6 +83,22 @@ impl ScanCache { Arc::clone(self.0.ranges.lock().get_or_insert(ranges)) } + /// The serialized bytes a whole-file scan of `projection` reads, computing them with + /// `compute` on first use. + pub(super) fn projection_bytes( + &self, + projection: &BoundExpression, + compute: impl FnOnce() -> VortexResult, + ) -> VortexResult { + let key = ExactBoundExpr(projection.clone()); + if let Some(bytes) = self.0.projection_bytes.lock().get(&key) { + return Ok(*bytes); + } + let bytes = compute()?; + self.0.projection_bytes.lock().insert(key, bytes); + Ok(bytes) + } + pub(super) fn values(&self) -> ValueCache { self.0.values.clone() } diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index c6b359fa328..cbb65c5959f 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -11,8 +11,9 @@ //! Import this module's [`ScanBuilder`] to select V2 explicitly. It accepts a [`ScanFile`] //! directly and supports projection, filters, row ranges, and selections. Unsupported layouts return an //! error; limits are not part of this builder's API. Work is partitioned into fixed -//! chunks of up to 65,536 rows, or 1,048,576 rows for filtered scans, independently of the stored -//! layout's chunk boundaries. +//! chunks of up to 65,536 rows, or up to 1,048,576 rows for filtered scans, independently of the +//! stored layout's chunk boundaries. Filtered scans use smaller chunks when the projection reads +//! many bytes per row or the row range would give fewer chunks than threads. //! Full-file projections use smaller 8,192-row multiples when serialized bytes per row are high. //! //! A [`ScanFile`] carrying a [`ScanCache`] shares its lowered plan, segment index, and dictionary @@ -45,6 +46,7 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_io::request::IoService; use vortex_io::session::RuntimeSessionExt; +use vortex_mask::Mask; use vortex_scan::planning::planner::WorkScope; use vortex_scan::selection::Selection; use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; @@ -55,6 +57,7 @@ pub use self::cache::ScanCache; use crate::LayoutRef; use crate::plan::physical::lower; use crate::plan::physical::optimize; +use crate::plan::physical::pipeline::Scan; use crate::plan::physical::plan_row_idx_expression; use crate::scan::arrow::to_record_batch; use crate::scan::execution::FilterPlans; @@ -76,13 +79,40 @@ fn rows_per_chunk(row_count: u64, locations: &[SegmentLocation]) -> usize { if row_count <= MAX_ROWS_PER_CHUNK as u64 { return MAX_ROWS_PER_CHUNK; } - let bytes: u64 = locations + let blocks = blocks_per_target_bytes(total_bytes(locations), row_count); + blocks.clamp(1, 8) as usize * ROW_BLOCK_SIZE +} + +/// Rows per task of a filtered scan over `rows` rows of the file. +/// +/// Large chunks amortize filter evaluation over cheap projections. A projection that reads many +/// bytes per row, or a row range that would otherwise give fewer chunks than threads, gets smaller +/// chunks so the surviving rows still spread across the runtime's threads. +fn filtered_rows_per_chunk( + rows: u64, + row_count: u64, + projected_bytes: u64, + parallelism: usize, +) -> usize { + let max_blocks = (FILTERED_ROWS_PER_CHUNK / ROW_BLOCK_SIZE) as u64; + let min_blocks = (MAX_ROWS_PER_CHUNK / ROW_BLOCK_SIZE) as u64; + let by_bytes = blocks_per_target_bytes(projected_bytes, row_count).max(1); + let by_threads = (rows / parallelism.max(1) as u64 / ROW_BLOCK_SIZE as u64).max(min_blocks); + let blocks = by_bytes.min(by_threads).min(max_blocks); + usize::try_from(blocks).vortex_expect("at most FILTERED_ROWS_PER_CHUNK blocks") * ROW_BLOCK_SIZE +} + +/// The row blocks that hold [`TARGET_BYTES_PER_CHUNK`] when `bytes` spread over `row_count` rows. +fn blocks_per_target_bytes(bytes: u64, row_count: u64) -> u64 { + let bytes_per_row = bytes.div_ceil(row_count.max(1)).max(1); + TARGET_BYTES_PER_CHUNK / bytes_per_row / ROW_BLOCK_SIZE as u64 +} + +fn total_bytes(locations: &[SegmentLocation]) -> u64 { + locations .iter() .map(|location| u64::from(location.length)) - .sum(); - let bytes_per_row = bytes.div_ceil(row_count).max(1); - let blocks = TARGET_BYTES_PER_CHUNK / bytes_per_row / ROW_BLOCK_SIZE as u64; - blocks.clamp(1, 8) as usize * ROW_BLOCK_SIZE + .sum() } /// Fixed boundaries in the file's row domain, clipped to the requested range. @@ -307,18 +337,40 @@ impl PreparedScan { None => { let plan = optimize(plan_row_idx_expression(builder.projection.clone(), root)?)?; - cache.insert_projection(builder.projection, plan) + cache.insert_projection(builder.projection.clone(), plan) } } }; + let rows = builder.row_range.unwrap_or(0..row_count); let chunk_rows = if filter.is_some() { - FILTERED_ROWS_PER_CHUNK + let projected_bytes = if identity { + total_bytes(&file.locations) + } else { + cache.projection_bytes(&builder.projection, || { + let segments = Scan::segments( + &builder.session, + &projection, + 0..row_count, + &Mask::new_true(usize::try_from(row_count)?), + )?; + Ok(segments + .iter() + .map(|segment| u64::from(file.locations[**segment as usize].length)) + .sum()) + })? + }; + filtered_rows_per_chunk( + rows.end - rows.start, + row_count, + projected_bytes, + get_available_parallelism().unwrap_or(1), + ) } else if identity { rows_per_chunk(row_count, &file.locations) } else { MAX_ROWS_PER_CHUNK }; - let splits = row_chunks(builder.row_range.unwrap_or(0..row_count), chunk_rows); + let splits = row_chunks(rows, chunk_rows); let ranges = match cache.ranges() { Some(ranges) => ranges, None => cache.insert_ranges(segment_ranges(&file.locations)),