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 all 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
140 changes: 104 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,47 @@ impl Not for &BooleanArray {
impl DaftLogical<&Self> for BooleanArray {
type Output = DaftResult<Self>;
fn and(&self, rhs: &Self) -> Self::Output {
// When performing a logical AND with a NULL value:
// - If the non-null value is false, the result is false (not null)
// - If the non-null value is true, the result is null
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(),
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 +651,18 @@ impl DaftLogical<&Self> for BooleanArray {
),
)))
}
(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))
}
}
(l, r) => Err(DaftError::ValueError(format!(
Expand All @@ -640,13 +674,47 @@ impl DaftLogical<&Self> for BooleanArray {
}

fn or(&self, rhs: &Self) -> Self::Output {
// When performing a logical OR with a NULL value:
// - If the non-null value is false, the result is null
// - If the non-null value is true, the result is true (not 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(),
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)),
(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 +724,18 @@ impl DaftLogical<&Self> for BooleanArray {
),
)))
}
(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))
}
}
(l, r) => Err(DaftError::ValueError(format!(
Expand Down Expand Up @@ -756,33 +824,33 @@ impl DaftCompare<&Self> for NullArray {
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