diff --git a/vortex-python/src/arrays/mod.rs b/vortex-python/src/arrays/mod.rs index 834b42ef718..06e94f8d464 100644 --- a/vortex-python/src/arrays/mod.rs +++ b/vortex-python/src/arrays/mod.rs @@ -48,7 +48,6 @@ use vortex::array::session::ArraySessionExt; use vortex::buffer::ByteBuffer; use vortex::dtype::DType; use vortex::dtype::Nullability; -use vortex::dtype::PType; use vortex::flatbuffers::WriteFlatBufferExt; use vortex::ipc::messages::EncoderMessage; use vortex::ipc::messages::MessageEncoder; @@ -415,11 +414,7 @@ impl PyArray { }; (*ptype, dtype) } else { - let ptype = if start > 0 && stop > 0 { - PType::U64 - } else { - PType::I64 - }; + let ptype = range_to_sequence::range_ptype(start, stop, step); let dtype = DType::Primitive(ptype, Nullability::NonNullable); (ptype, dtype) }; diff --git a/vortex-python/src/arrays/range_to_sequence.rs b/vortex-python/src/arrays/range_to_sequence.rs index eed03c64d13..bae9227ebdb 100644 --- a/vortex-python/src/arrays/range_to_sequence.rs +++ b/vortex-python/src/arrays/range_to_sequence.rs @@ -9,6 +9,7 @@ use vortex::buffer::Buffer; use vortex::dtype::DType; use vortex::dtype::NativePType; use vortex::dtype::Nullability; +use vortex::dtype::PType; use vortex::encodings::sequence::Sequence; use vortex::error::VortexExpect; use vortex::error::VortexResult; @@ -25,13 +26,17 @@ pub fn sequence_array_from_range + Into> vortex_bail!("Step must not be zero"); } - let Some(len) = range_len(start, stop, step) else { + // A `Sequence` holds at least one element, so an empty range returns a primitive array + // instead. `range_len` reports a range that runs the wrong way as `None` and one whose bounds + // coincide as `Some(0)`; Python's `range` calls both empty. + let len = range_len(start, stop, step).unwrap_or(0); + if len == 0 { let validity = match dtype.nullability() { Nullability::NonNullable => Validity::NonNullable, Nullability::Nullable => Validity::AllValid, }; return Ok(PrimitiveArray::new::(Buffer::empty(), validity).into_array()); - }; + } let Ok(start) = T::try_from(start) else { vortex_bail!( "Start, {}, does not fit in requested dtype: {}", @@ -46,6 +51,18 @@ pub fn sequence_array_from_range + Into> Ok(Sequence::try_new_typed::(start, step, dtype.nullability(), len)?.into_array()) } +/// The [`PType`] a Python `range` converts to when the caller does not request a dtype. +/// +/// An unsigned type needs a positive step as well as positive bounds, because the step is stored +/// in the same type as the values. +pub fn range_ptype(start: isize, stop: isize, step: isize) -> PType { + if start > 0 && stop > 0 && step > 0 { + PType::U64 + } else { + PType::I64 + } +} + fn range_len(start: isize, stop: isize, step: isize) -> Option { if step > 0 { if start > stop { @@ -72,18 +89,69 @@ fn range_len(start: isize, stop: isize, step: isize) -> Option { #[cfg(test)] mod test { + use rstest::rstest; use vortex::array::IntoArray as _; use vortex::array::assert_arrays_eq; + use vortex::array::match_each_integer_ptype; use vortex::buffer::buffer; use vortex::dtype::DType; use vortex::dtype::Nullability; use vortex::dtype::PType; + use vortex::error::VortexResult; use vortex_array::VortexSessionExecute; use vortex_array::array_session; use crate::arrays::range_to_sequence::range_len; + use crate::arrays::range_to_sequence::range_ptype; use crate::arrays::range_to_sequence::sequence_array_from_range; + /// Python's `range` is empty in three shapes, and `range_len` reports the first as `Some(0)` + /// and the other two as `None`. All three must convert to an empty array. + #[rstest] + #[case::bounds_coincide(3, 3, 1)] + #[case::bounds_coincide_below_zero(-5, -5, 1)] + #[case::bounds_coincide_with_negative_step(3, 3, -1)] + #[case::positive_step_runs_backwards(10, 3, 1)] + #[case::negative_step_runs_forwards(0, 10, -1)] + fn empty_range_converts_to_an_empty_array( + #[case] start: isize, + #[case] stop: isize, + #[case] step: isize, + ) -> VortexResult<()> { + let dtype = DType::Primitive(PType::I64, Nullability::NonNullable); + let arr = sequence_array_from_range::(start, stop, step, dtype)?; + assert_eq!(arr.len(), 0); + Ok(()) + } + + /// A descending range needs a signed type, because the negative step is stored in the same + /// type as the values. + #[rstest] + #[case::ascending_above_zero(1, 10, 1, PType::U64)] + #[case::ascending_from_zero(0, 10, 1, PType::I64)] + #[case::descending_above_zero(5, 1, -1, PType::I64)] + #[case::descending_through_zero(5, -1, -1, PType::I64)] + fn range_ptype_holds_the_step( + #[case] start: isize, + #[case] stop: isize, + #[case] step: isize, + #[case] expected: PType, + ) { + assert_eq!(range_ptype(start, stop, step), expected); + } + + #[test] + fn descending_range_above_zero_converts() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let ptype = range_ptype(5, 1, -1); + let dtype = DType::Primitive(ptype, Nullability::NonNullable); + let arr = match_each_integer_ptype!(ptype, |T| { + sequence_array_from_range::(5, 1, -1, dtype) + })?; + assert_arrays_eq!(arr, buffer![5i64, 4, 3, 2].into_array(), &mut ctx); + Ok(()) + } + #[test] fn test_range_len() { assert_eq!(range_len(0, 10, 1).unwrap(), 10); diff --git a/vortex-python/test/test_from_range.py b/vortex-python/test/test_from_range.py index 15de76db734..13f13551227 100644 --- a/vortex-python/test/test_from_range.py +++ b/vortex-python/test/test_from_range.py @@ -49,3 +49,13 @@ def test_from_range_invalid(): arr = vx.array(range(0, 10, -1)) assert values(arr) == [] + + +def test_from_range_3_3(): + arr = vx.array(range(3, 3)) + assert values(arr) == [] + + +def test_from_range_5_1_minus_1(): + arr = vx.array(range(5, 1, -1)) + assert values(arr) == list(range(5, 1, -1))