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.

Sign up for GitHub

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

fix: boolean and/or expressions with null #3544

Merged
merged 4 commits into from
Dec 11, 2024
Merged

fix: boolean and/or expressions with null #3544

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
ac52f08
fix: boolean and/or expressions with null
kevinzwang Dec 10, 2024
3fc0ce8
fix another test
kevinzwang Dec 11, 2024
e0468af
fix last test
kevinzwang Dec 11, 2024
54f9e66
update comments
kevinzwang Dec 11, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
136 changes: 100 additions & 36 deletions src/daft-core/src/array/ops/comparison.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -601,13 +601,45 @@
impl DaftLogical<&Self> for BooleanArray {
type Output = DaftResult<Self>;
fn and(&self, rhs: &Self) -> Self::Output {
// keep all valid false values when and with null
kevinzwang marked this conversation as resolved.
Show resolved Hide resolved
fn and_with_null(name: &str, arr: &BooleanArray) -> BooleanArray {
let values = arr.as_arrow().values();

let new_validity = match arr.as_arrow().validity() {
None => values.not(),

Check warning on line 609 in src/daft-core/src/array/ops/comparison.rs

View check run for this annotation

Codecov / codecov/patch

src/daft-core/src/array/ops/comparison.rs#L609 

Added line #L609 was not covered by tests
Some(validity) => arrow2::bitmap::and(&values.not(), validity),
};

BooleanArray::from((
name,
arrow2::array::BooleanArray::new(
arrow2::datatypes::DataType::Boolean,
values.clone(),
Some(new_validity),
),
))
}

match (self.len(), rhs.len()) {
(x, y) if x == y => {
Copy link
Contributor

Choose a reason for hiding this comment

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

maybe better names than x and y?

Copy link
Member Author

@kevinzwang kevinzwang Dec 11, 2024 •
edited
Loading

Choose a reason for hiding this comment

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

This naming seems to exist across all of the ops so I think I will just leave it be. Don't want to make this PR too big or stray away from convention

let validity =
arrow_bitmap_and_helper(self.as_arrow().validity(), rhs.as_arrow().validity());

let result_bitmap =
arrow2::bitmap::and(self.as_arrow().values(), rhs.as_arrow().values());
let l_values = self.as_arrow().values();
let r_values = rhs.as_arrow().values();

// (false & NULL) should be false, compute validity to ensure that
let validity = match (self.as_arrow().validity(), rhs.as_arrow().validity()) {
(None, None) => None,
(None, Some(r_valid)) => Some(arrow2::bitmap::or(&l_values.not(), r_valid)),
(Some(l_valid), None) => Some(arrow2::bitmap::or(l_valid, &r_values.not())),
(Some(l_valid), Some(r_valid)) => Some(arrow2::bitmap::or(
&arrow2::bitmap::or(
&arrow2::bitmap::and(&l_values.not(), l_valid),
&arrow2::bitmap::and(&r_values.not(), r_valid),
),
&arrow2::bitmap::and(l_valid, r_valid),
)),
};

let result_bitmap = arrow2::bitmap::and(l_values, r_values);
Ok(Self::from((
self.name(),
arrow2::array::BooleanArray::new(
Expand All @@ -617,18 +649,18 @@
),
)))
}
(l_size, 1) => {
(_, 1) => {
andrewgazelka marked this conversation as resolved.
Show resolved Hide resolved
if let Some(value) = rhs.get(0) {
self.and(value)
} else {
Ok(Self::full_null(self.name(), &DataType::Boolean, l_size))
Ok(and_with_null(self.name(), self))
}
}
(1, r_size) => {
(1, _) => {
if let Some(value) = self.get(0) {
rhs.and(value)
} else {
Ok(Self::full_null(self.name(), &DataType::Boolean, r_size))
Ok(and_with_null(self.name(), rhs))

Check warning on line 663 in src/daft-core/src/array/ops/comparison.rs

View check run for this annotation

Codecov / codecov/patch

src/daft-core/src/array/ops/comparison.rs#L663 

Added line #L663 was not covered by tests
}
}
(l, r) => Err(DaftError::ValueError(format!(
Expand All @@ -640,13 +672,45 @@
}

fn or(&self, rhs: &Self) -> Self::Output {
// keep all valid true values when or with null
fn or_with_null(name: &str, arr: &BooleanArray) -> BooleanArray {
let values = arr.as_arrow().values();

let new_validity = match arr.as_arrow().validity() {
None => values.clone(),

Check warning on line 680 in src/daft-core/src/array/ops/comparison.rs

View check run for this annotation

Codecov / codecov/patch

src/daft-core/src/array/ops/comparison.rs#L680 

Added line #L680 was not covered by tests
Some(validity) => arrow2::bitmap::and(values, validity),
};

BooleanArray::from((
name,
arrow2::array::BooleanArray::new(
arrow2::datatypes::DataType::Boolean,
values.clone(),
Some(new_validity),
),
))
}

match (self.len(), rhs.len()) {
(x, y) if x == y => {
andrewgazelka marked this conversation as resolved.
Show resolved Hide resolved
let validity =
arrow_bitmap_and_helper(self.as_arrow().validity(), rhs.as_arrow().validity());

let result_bitmap =
arrow2::bitmap::or(self.as_arrow().values(), rhs.as_arrow().values());
let l_values = self.as_arrow().values();
let r_values = rhs.as_arrow().values();

// (true | NULL) should be true, compute validity to ensure that
let validity = match (self.as_arrow().validity(), rhs.as_arrow().validity()) {
(None, None) => None,
(None, Some(r_valid)) => Some(arrow2::bitmap::or(l_values, r_valid)),
(Some(l_valid), None) => Some(arrow2::bitmap::or(l_valid, r_values)),

Check warning on line 703 in src/daft-core/src/array/ops/comparison.rs

View check run for this annotation

Codecov / codecov/patch

src/daft-core/src/array/ops/comparison.rs#L702-L703 

Added lines #L702 - L703 were not covered by tests
(Some(l_valid), Some(r_valid)) => Some(arrow2::bitmap::or(
&arrow2::bitmap::or(
&arrow2::bitmap::and(l_values, l_valid),
&arrow2::bitmap::and(r_values, r_valid),
),
&arrow2::bitmap::and(l_valid, r_valid),
)),
};

let result_bitmap = arrow2::bitmap::or(l_values, r_values);
Ok(Self::from((
self.name(),
arrow2::array::BooleanArray::new(
Expand All @@ -656,18 +720,18 @@
),
)))
}
(l_size, 1) => {
(_, 1) => {
if let Some(value) = rhs.get(0) {
self.or(value)
} else {
Ok(Self::full_null(self.name(), &DataType::Boolean, l_size))
Ok(or_with_null(self.name(), self))
}
}
(1, r_size) => {
(1, _) => {
if let Some(value) = self.get(0) {
rhs.or(value)
} else {
Ok(Self::full_null(self.name(), &DataType::Boolean, r_size))
Ok(or_with_null(self.name(), rhs))

Check warning on line 734 in src/daft-core/src/array/ops/comparison.rs

View check run for this annotation

Codecov / codecov/patch

src/daft-core/src/array/ops/comparison.rs#L734 

Added line #L734 was not covered by tests
}
}
(l, r) => Err(DaftError::ValueError(format!(
Expand Down Expand Up @@ -756,33 +820,33 @@
impl DaftLogical<bool> for BooleanArray {
type Output = DaftResult<Self>;
fn and(&self, rhs: bool) -> Self::Output {
let validity = self.as_arrow().validity();
if rhs {
Ok(self.clone())
} else {
use arrow2::{array, bitmap::Bitmap, datatypes::DataType};
let arrow_array = array::BooleanArray::new(
DataType::Boolean,
Bitmap::new_zeroed(self.len()),
validity.cloned(),
);
Ok(Self::from((self.name(), arrow_array)))
Ok(Self::from((
self.name(),
arrow2::array::BooleanArray::new(
arrow2::datatypes::DataType::Boolean,
arrow2::bitmap::Bitmap::new_zeroed(self.len()),
None, // false & x is always valid false for any x
),
)))
}
}

fn or(&self, rhs: bool) -> Self::Output {
let validity = self.as_arrow().validity();
if rhs {
use arrow2::{array, bitmap::Bitmap, datatypes::DataType};
let arrow_array = array::BooleanArray::new(
DataType::Boolean,
Bitmap::new_zeroed(self.len()).not(),
validity.cloned(),
);
return Ok(Self::from((self.name(), arrow_array)));
Ok(Self::from((
self.name(),
arrow2::array::BooleanArray::new(
arrow2::datatypes::DataType::Boolean,
arrow2::bitmap::Bitmap::new_trued(self.len()),
None, // true | x is always valid true for any x
),
)))
} else {
Ok(self.clone())
}

Ok(self.clone())
}

fn xor(&self, rhs: bool) -> Self::Output {
Expand Down
Loading
Loading