[FEAT] List slice expression #2479

desmondcheongzx · 2024-07-04T01:50:26Z

Adds an expression that returns a specified subset of elements from a list.

The API to use this expression is daft.Expression.list.slice(start: int | Expression, end: int | Expression).

start is the 0-indexed position of the list to start retrieving elements. If this value is negative, then it's the position from the end of the array. i.e. -1 points to the last element of the array
end is the 0-indexed position of the list to stop retrieving elements. If this value is negative, then it's the position from the end of the array. i.e. -1 points to the last element of the array

daft/expressions/expressions.py

jaychia · 2024-07-04T02:39:44Z

daft/series.py

@@ -541,6 +541,8 @@ def sum(self) -> Series:
        assert self._series is not None
        return Series._from_pyseries(self._series.sum())

+    # TODO: Is this needed?


(Mostly for testing only)

jaychia · 2024-07-04T02:44:36Z

src/daft-core/src/array/ops/list.rs

+    // Given an i64 array that may have either 1 or `self.len()` elements, create an iterator with
+    // `self.len()` elements. If there was originally 1 element, we repeat this element `self.len()`
+    // times, otherwise we simply take the original array.
+    fn create_iter<'a>(&'a self, arr: &'a Int64Array) -> Box<dyn Iterator<Item = i64> + '_> {


Let's dedup this

jaychia · 2024-07-04T02:47:49Z

src/daft-core/src/array/ops/list.rs

            }
        }
+        let total_capacity = *new_offsets.last().unwrap();
+        let default = Series::full_null("null", self.child_data_type(), 1);


I think our growable allows you to add_null without needing a dummy array

Removed this unnecessary default value

jaychia · 2024-07-04T02:50:46Z

src/daft-core/src/array/ops/list.rs

+            self.name(),
+            self.child_data_type(),
+            vec![&self.flat_child, &default],
+            true, // TODO is this right?


If you are very sure that you won't be adding any nulls, you can set this to false for some performance gains.

jaychia · 2024-07-04T02:52:51Z

src/daft-core/src/array/ops/list.rs

+            total_capacity as usize,
+        );
+        for (i, start) in slicing_indexes.iter().enumerate() {
+            if *start >= 0 {


Let's follow python behavior instead?

Sounds good, changed the API

jaychia · 2024-07-04T02:54:50Z

src/daft-core/src/array/ops/list.rs

+        start_iter: impl Iterator<Item = i64>,
+        length_iter: impl Iterator<Item = i64>,
+    ) -> DaftResult<Series> {
+        let list_size = self.fixed_element_len();


We can probably dedup a lot of the logic, just need to turn the offsets into an interable or something.

jaychia · 2024-07-04T02:55:38Z

src/daft-core/src/series/ops/list.rs

+                self.fixed_size_list()?.get_slices(start_arr, length_arr)
+            }
+            dt => Err(DaftError::TypeError(format!(
+                "list splice not implemented for {}",


Suggested change

"list splice not implemented for {}",

"list slice not implemented for {}",

Thanks for the catch! Fixed

jaychia · 2024-07-04T03:00:26Z

tests/table/list/test_list_slice.py

+        table.eval_expression_list([col("col").list.slice()])
+    with pytest.raises(TypeError, match="takes 3 positional arguments but 4 were given"):
+        table.eval_expression_list([col("col").list.slice(0, 0, 0)])
+


Also other corner-case (completely empty MicroPartition). This comes up quite a lot, especially in distributed settings when a filter is applied and completely removes data from a MicroPartition (making it empty).

You can make a well-typed empty series like this:

pa.array([], type=pa.list_(...))

Added this test

jaychia · 2024-07-04T03:01:41Z

tests/table/list/test_list_slice.py

+def test_list_slice_non_list_type():
+    table = MicroPartition.from_pydict(
+        {
+            "mapcol": [{"a": 1}, {"b": 1}, {"c": 1}],


This is a struct 😂

desmondcheongzx

@jaychia thanks for the review! I addressed your comments

desmondcheongzx · 2024-07-05T21:11:19Z

tests/table/list/test_list_slice.py

+        table.eval_expression_list([col("col").list.slice()])
+    with pytest.raises(TypeError, match="takes 3 positional arguments but 4 were given"):
+        table.eval_expression_list([col("col").list.slice(0, 0, 0)])
+


Added this test

daft/expressions/expressions.py

desmondcheongzx · 2024-07-05T21:12:18Z

daft/series.py

@@ -541,6 +541,8 @@ def sum(self) -> Series:
        assert self._series is not None
        return Series._from_pyseries(self._series.sum())

+    # TODO: Is this needed?


desmondcheongzx · 2024-07-05T21:12:25Z

src/daft-core/src/array/ops/list.rs

+    // Given an i64 array that may have either 1 or `self.len()` elements, create an iterator with
+    // `self.len()` elements. If there was originally 1 element, we repeat this element `self.len()`
+    // times, otherwise we simply take the original array.
+    fn create_iter<'a>(&'a self, arr: &'a Int64Array) -> Box<dyn Iterator<Item = i64> + '_> {


desmondcheongzx · 2024-07-05T21:13:04Z

src/daft-core/src/array/ops/list.rs

            }
        }
+        let total_capacity = *new_offsets.last().unwrap();
+        let default = Series::full_null("null", self.child_data_type(), 1);


Removed this unnecessary default value

desmondcheongzx · 2024-07-05T21:13:27Z

src/daft-core/src/array/ops/list.rs

+            self.name(),
+            self.child_data_type(),
+            vec![&self.flat_child, &default],
+            true, // TODO is this right?


desmondcheongzx · 2024-07-05T21:13:37Z

src/daft-core/src/array/ops/list.rs

+            total_capacity as usize,
+        );
+        for (i, start) in slicing_indexes.iter().enumerate() {
+            if *start >= 0 {


Sounds good, changed the API

desmondcheongzx · 2024-07-05T21:13:47Z

src/daft-core/src/array/ops/list.rs

+        start_iter: impl Iterator<Item = i64>,
+        length_iter: impl Iterator<Item = i64>,
+    ) -> DaftResult<Series> {
+        let list_size = self.fixed_element_len();


desmondcheongzx · 2024-07-05T21:14:00Z

src/daft-core/src/series/ops/list.rs

+                self.fixed_size_list()?.get_slices(start_arr, length_arr)
+            }
+            dt => Err(DaftError::TypeError(format!(
+                "list splice not implemented for {}",


Thanks for the catch! Fixed

desmondcheongzx · 2024-07-05T21:14:18Z

tests/table/list/test_list_slice.py

+def test_list_slice_non_list_type():
+    table = MicroPartition.from_pydict(
+        {
+            "mapcol": [{"a": 1}, {"b": 1}, {"c": 1}],


jaychia

Lookin good -- congrats on 1st PR!!!!!!!!!!!!

jaychia · 2024-07-05T23:28:37Z

src/daft-core/src/array/ops/list.rs

+}
+
+pub fn get_slices_helper(
+    parent_offsets: &arrow2::offset::OffsetsBuffer<i64>,


You can actually make this an impl Iterator<Item = i64> (since you mostly use it that way anyways!) and then the fixed size list case would be quite clean: std::iter::repeat(element_len).take(list_len)

Done, thanks for the suggestion!

github-actions bot added the enhancement New feature or request label Jul 4, 2024

desmondcheongzx force-pushed the add-list-slice branch from f71d007 to 1099dd0 Compare July 4, 2024 02:06

Initial implementation

446ee32

desmondcheongzx force-pushed the add-list-slice branch from 1099dd0 to 446ee32 Compare July 4, 2024 02:07

jaychia reviewed Jul 4, 2024

View reviewed changes

Address comments

f5c5019

desmondcheongzx commented Jul 5, 2024

View reviewed changes

jaychia approved these changes Jul 5, 2024

View reviewed changes

Address comment

67d4ec4

desmondcheongzx enabled auto-merge (squash) July 6, 2024 00:32

desmondcheongzx merged commit 7b81de6 into Eventual-Inc:main Jul 6, 2024
41 checks passed

desmondcheongzx mentioned this pull request Jul 9, 2024

Add slice expression #2478

Closed

desmondcheongzx deleted the add-list-slice branch August 1, 2024 04:41

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[FEAT] List slice expression #2479

[FEAT] List slice expression #2479

desmondcheongzx commented Jul 4, 2024 •

edited

Loading

jaychia Jul 4, 2024

desmondcheongzx Jul 5, 2024

jaychia Jul 4, 2024

desmondcheongzx Jul 5, 2024

jaychia Jul 4, 2024

desmondcheongzx Jul 5, 2024

jaychia Jul 4, 2024

desmondcheongzx Jul 5, 2024

jaychia Jul 4, 2024

desmondcheongzx Jul 5, 2024

jaychia Jul 4, 2024

desmondcheongzx Jul 5, 2024

jaychia Jul 4, 2024

desmondcheongzx Jul 5, 2024

jaychia Jul 4, 2024

desmondcheongzx Jul 5, 2024

jaychia Jul 4, 2024

desmondcheongzx Jul 5, 2024

desmondcheongzx left a comment

desmondcheongzx Jul 5, 2024

desmondcheongzx Jul 5, 2024

desmondcheongzx Jul 5, 2024

desmondcheongzx Jul 5, 2024

desmondcheongzx Jul 5, 2024

desmondcheongzx Jul 5, 2024

desmondcheongzx Jul 5, 2024

desmondcheongzx Jul 5, 2024

desmondcheongzx Jul 5, 2024

jaychia left a comment

jaychia Jul 5, 2024

desmondcheongzx Jul 6, 2024

	"list splice not implemented for {}",
	"list slice not implemented for {}",

[FEAT] List slice expression #2479

[FEAT] List slice expression #2479

Conversation

desmondcheongzx commented Jul 4, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

desmondcheongzx left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jaychia left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

desmondcheongzx commented Jul 4, 2024 •

edited

Loading