Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[FEAT] Add getter for Struct and List expressions #1775

Merged
merged 9 commits into from
Jan 12, 2024

Conversation

kevinzwang
Copy link
Member

@kevinzwang kevinzwang commented Jan 10, 2024

This PR adds getters for both lists and structs, which enables the access of elements in ListArray, FixedSizeListArray, and StructArray types.

For lists:

Expression.list.get(idx, default=None)

This method takes in an index, which can be a single value or an expression, and optionally a default value, which will be set when the index is out of bounds. It retrieves the value at index of each list in the array.

For structs:

Expression.struct.get(name)

This method takes in the name of the field to get, and returns an expression of that struct's field.

@github-actions github-actions bot added the enhancement New feature or request label Jan 10, 2024
@kevinzwang kevinzwang linked an issue Jan 10, 2024 that may be closed by this pull request
Copy link

codecov bot commented Jan 10, 2024

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (a994c7b) 84.60% compared to head (e945aca) 84.88%.
Report is 6 commits behind head on main.

❗ Current head e945aca differs from pull request most recent head 47fa923. Consider uploading reports for the commit 47fa923 to get more accurate results

Additional details and impacted files

Impacted file tree graph

@@            Coverage Diff             @@
##             main    #1775      +/-   ##
==========================================
+ Coverage   84.60%   84.88%   +0.28%     
==========================================
  Files          55       55              
  Lines        5554     5604      +50     
==========================================
+ Hits         4699     4757      +58     
+ Misses        855      847       -8     
Files Coverage Δ
daft/expressions/expressions.py 92.37% <100.00%> (+0.28%) ⬆️
daft/series.py 92.50% <100.00%> (-0.13%) ⬇️

... and 8 files with indirect coverage changes

idx_iter: &mut impl Iterator<Item = Option<i64>>,
default: &Series,
) -> DaftResult<Series> {
let default = default.cast(self.child_data_type())?;
Copy link
Member Author

@kevinzwang kevinzwang Jan 10, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

default isn't a reserved keyword in Rust but the syntax highlighting is different on it. Maybe because it's in std.

Copy link
Contributor

@jaychia jaychia left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looking good so far!

src/daft-core/src/array/ops/list.rs Outdated Show resolved Hide resolved
src/daft-core/src/array/ops/list.rs Outdated Show resolved Hide resolved
src/daft-core/src/series/ops/list.rs Outdated Show resolved Hide resolved
src/daft-dsl/src/functions/list/get.rs Show resolved Hide resolved
);

// set all default indices to 0 if there is only one
let multiple_defaults = match default.len() {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we can probably make an assertion that default has length 1 and only allow for that case from the user-API.

Should be pretty rare that a user wants to have a different default per-element!

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In Python-land you can wrap the default into a LiteralExpression like so:

def get(idx: int, default: Any):
    default_expr = lit(default)
    ...

Then when it gets evaluated in Rust it should evaluate the a Series with length 1 I believe (we should double-check!)

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks! This method does seem to work, although it doesn't seem to emit any errors when you pass in an Expression, it just manages it as a generic PyObject.

I think this is probably better handled on the Rust side, but is there a way to check of an expression's type is castable to another's?

@kevinzwang kevinzwang marked this pull request as ready for review January 11, 2024 21:48
Copy link
Contributor

@jaychia jaychia left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looking good :)

daft/expressions/expressions.py Outdated Show resolved Hide resolved
src/daft-core/src/array/ops/list.rs Show resolved Hide resolved
src/daft-core/src/array/ops/list.rs Show resolved Hide resolved
src/daft-dsl/src/functions/struct_/field.rs Outdated Show resolved Hide resolved
tests/table/list/test_list_get.py Show resolved Hide resolved
tests/table/struct/test_struct_field.py Outdated Show resolved Hide resolved
@kevinzwang
Copy link
Member Author

Sorry that this PR is getting a little larger than expected now, most of the changed files are just adding implementations for with_validity() and validity() which are relatively straightforward.

Copy link
Contributor

@jaychia jaychia left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great work!

Let's also add these expressions into our documentation: https://github.com/Eventual-Inc/Daft/blob/main/docs/source/api_docs/expressions.rst#nested

src/daft-dsl/src/functions/struct_/get.rs Outdated Show resolved Hide resolved
src/daft-dsl/src/functions/struct_/get.rs Outdated Show resolved Hide resolved
@github-actions github-actions bot added the documentation Improvements or additions to documentation label Jan 12, 2024
@kevinzwang kevinzwang enabled auto-merge (squash) January 12, 2024 23:22
@kevinzwang kevinzwang merged commit 1115869 into main Jan 12, 2024
38 checks passed
@kevinzwang kevinzwang deleted the kevin/struct-list-getters branch January 12, 2024 23:30
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
documentation Improvements or additions to documentation enhancement New feature or request
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Add getter for Struct and List expressions
2 participants