fix: fix reverse variable row decoding (pola-rs#13587)

meghantorchia · Jan 10, 2024 · 439de72 · 439de72
1 parent 2b43fc1
commit 439de72
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 1 deletion.
diff --git a/crates/polars-row/src/encode.rs b/crates/polars-row/src/encode.rs
@@ -307,6 +307,7 @@ mod test {
     use arrow::array::{Int32Array, Utf8Array};
 
     use super::*;
+    use crate::decode::decode_rows_from_binary;
     use crate::variable::{decode_binary, BLOCK_SIZE, EMPTY_SENTINEL, NON_EMPTY_SENTINEL};
 
     #[test]
@@ -392,4 +393,25 @@ mod test {
         let decoded = unsafe { decode_binary(&mut rows, &field) };
         assert_eq!(decoded, arr);
     }
+
+    #[test]
+    fn test_reverse_variable() {
+        let a = Utf8Array::<i64>::from_slice(["one", "two", "three", "four", "five", "six"]);
+
+        let fields = &[SortField {
+            descending: true,
+            nulls_last: false,
+        }];
+
+        let dtypes = [ArrowDataType::LargeUtf8];
+
+        unsafe {
+            let encoded = convert_columns(&[Box::new(a.clone())], fields);
+            let out = decode_rows_from_binary(&encoded.into_array(), fields, &dtypes, &mut vec![]);
+
+            let arr = &out[0];
+            let decoded = arr.as_any().downcast_ref::<Utf8Array<i64>>().unwrap();
+            assert_eq!(decoded, &a);
+        }
+    }
 }
diff --git a/crates/polars-row/src/variable.rs b/crates/polars-row/src/variable.rs
@@ -239,6 +239,7 @@ pub(super) unsafe fn decode_binary(rows: &mut [&[u8]], field: &SortField) -> Bin
             continuation_token,
             field.descending,
         );
+        let values_offset = values.len();
 
         let mut to_read = str_len;
         // we start at one, as we skip the validity byte
@@ -258,7 +259,10 @@ pub(super) unsafe fn decode_binary(rows: &mut [&[u8]], field: &SortField) -> Bin
         offsets.push(values.len() as i64);
 
         if field.descending {
-            values.iter_mut().for_each(|o| *o = !*o)
+            values
+                .get_unchecked_release_mut(values_offset..)
+                .iter_mut()
+                .for_each(|o| *o = !*o)
         }
     }
 

diff --git a/py-polars/tests/unit/streaming/test_streaming_sort.py b/py-polars/tests/unit/streaming/test_streaming_sort.py
@@ -232,3 +232,15 @@ def test_streaming_sort_fixed_reverse() -> None:
     assert_df_sorted_by(
         df, q.collect(streaming=False), ["a", "b"], descending=descending
     )
+
+
+def test_reverse_variable_sort_13573() -> None:
+    df = pl.DataFrame(
+        {
+            "a": ["one", "two", "three"],
+            "b": ["four", "five", "six"],
+        }
+    ).lazy()
+    assert df.sort("a", "b", descending=[True, False]).collect(streaming=True).to_dict(
+        as_series=False
+    ) == {"a": ["two", "three", "one"], "b": ["five", "six", "four"]}