Skip to content

ARROW-4605: [Rust] Move filter and limit code from DataFusion into compute module #3741

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

Closed
nvizzari wants to merge 14 commits into from
Closed

ARROW-4605: [Rust] Move filter and limit code from DataFusion into compute module #3741

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
d216fa0
Move filter and limit to array_ops
Feb 24, 2019
0ca0412
Rewrite filter and limit using macros
Feb 24, 2019
32a2f85
Add tests for limit and filter
Feb 24, 2019
b20ea6d
Do bound checking in limit function
Feb 24, 2019
2a389a3
create BinaryArray directly from byte slice to prevent converting to …
nevi-me Feb 24, 2019
6422e18
cargo fmt
nevi-me Feb 24, 2019
2f44a8a
Use the size of the array as limit instead of returning an error
Feb 24, 2019
2e9616b
Add documentation for the limit function
Feb 24, 2019
5a1047c
Name variables consistently
Feb 24, 2019
58d1f5c
Merge branch 'ARROW-4605' into ARROW-4605
Feb 24, 2019
728884b
Merge pull request #1 from nevi-me/ARROW-4605
Feb 24, 2019
f0578f6
Add tests for limit and filter with BinaryArray
Feb 24, 2019
257d235
Add support for null values in limit and filter
Feb 24, 2019
344379a
Initialize vectors with a capacity
Feb 24, 2019
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
137 changes: 135 additions & 2 deletions rust/arrow/src/compute/array_ops.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,9 +18,14 @@
//! Defines primitive computations on arrays, e.g. addition, equality, boolean logic.

use std::ops::Add;
use std::sync::Arc;

use crate::array::{Array, BooleanArray, PrimitiveArray};
use crate::datatypes::ArrowNumericType;
use crate::array::{
Array, ArrayRef, BinaryArray, BooleanArray, Float32Array, Float64Array, Int16Array,
Int32Array, Int64Array, Int8Array, PrimitiveArray, UInt16Array, UInt32Array,
UInt64Array, UInt8Array,
};
use crate::datatypes::{ArrowNumericType, DataType};
use crate::error::{ArrowError, Result};

/// Returns the minimum value in the array, according to the natural order.
Expand Down Expand Up @@ -204,6 +209,101 @@ where
Ok(b.finish())
}

macro_rules! filter_array {
($array:expr, $filter:expr, $array_type:ident) => {{
let b = $array.as_any().downcast_ref::<$array_type>().unwrap();
let mut builder = $array_type::builder(b.len());
for i in 0..b.len() {
if $filter.value(i) {
builder.append_value(b.value(i))?;
}
}
Ok(Arc::new(builder.finish()))
}};
}

pub fn filter(array: &Array, filter: &BooleanArray) -> Result<ArrayRef> {
match array.data_type() {
DataType::UInt8 => filter_array!(array, filter, UInt8Array),
DataType::UInt16 => filter_array!(array, filter, UInt16Array),
DataType::UInt32 => filter_array!(array, filter, UInt32Array),
DataType::UInt64 => filter_array!(array, filter, UInt64Array),
DataType::Int8 => filter_array!(array, filter, Int8Array),
DataType::Int16 => filter_array!(array, filter, Int16Array),
DataType::Int32 => filter_array!(array, filter, Int32Array),
DataType::Int64 => filter_array!(array, filter, Int64Array),
DataType::Float32 => filter_array!(array, filter, Float32Array),
DataType::Float64 => filter_array!(array, filter, Float64Array),
DataType::Boolean => filter_array!(array, filter, BooleanArray),
DataType::Utf8 => {
//TODO: this is inefficient and we should improve the Arrow impl to help make
// this more concise
let b = array.as_any().downcast_ref::<BinaryArray>().unwrap();
let mut values: Vec<String> = Vec::with_capacity(b.len());
for i in 0..b.len() {
if filter.value(i) {
values.push(b.get_string(i));
}
}
let tmp: Vec<&str> = values.iter().map(|s| s.as_str()).collect();
Ok(Arc::new(BinaryArray::from(tmp)))
}
other => Err(ArrowError::ComputeError(format!(
"filter not supported for {:?}",
other
))),
}
}

macro_rules! limit_array {
($array:expr, $num_rows_to_read:expr, $array_type:ident) => {{
let b = $array.as_any().downcast_ref::<$array_type>().unwrap();
let mut builder = $array_type::builder($num_rows_to_read);
for i in 0..$num_rows_to_read {
builder.append_value(b.value(i))?;
}
Ok(Arc::new(builder.finish()))
}};
}

pub fn limit(array: &Array, num_rows_to_read: usize) -> Result<ArrayRef> {
if array.len() < num_rows_to_read {

@nevi-me nevi-me Feb 24, 2019

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

I'm curious as to why we return an error here. Why not set the limit to the lower of the 2 instead?

vec![1,2,3,4,5].iter().take(7) doesn't throw an error, and SELECT * FROM table_with_100_records LIMIT 200 will also only return the 100 records. What do you think about the suggestion?

@nvizzari nvizzari Feb 24, 2019

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

My initial thought was that if someone calls this function with a num_rows_to_read larger than array.len(), they will probably want to know rather than silently choosing the minimum.

However given take and LIMIT behavior this is indeed probably not what is expected from this function. I will change it, good call.

@nevi-me nevi-me Feb 24, 2019

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

I think if you limit the num_rows_to_read to the array length.
Have a look at https://issues.apache.org/jira/browse/ARROW-3954, which will make filter more efficient in future.

return Err(ArrowError::ComputeError(format!(
"Number of rows to read ({:?}) is larger than the length of the array ({:?})",
num_rows_to_read,
array.len()
)));
}

match array.data_type() {
DataType::UInt8 => limit_array!(array, num_rows_to_read, UInt8Array),
DataType::UInt16 => limit_array!(array, num_rows_to_read, UInt16Array),
DataType::UInt32 => limit_array!(array, num_rows_to_read, UInt32Array),
DataType::UInt64 => limit_array!(array, num_rows_to_read, UInt64Array),
DataType::Int8 => limit_array!(array, num_rows_to_read, Int8Array),
DataType::Int16 => limit_array!(array, num_rows_to_read, Int16Array),
DataType::Int32 => limit_array!(array, num_rows_to_read, Int32Array),
DataType::Int64 => limit_array!(array, num_rows_to_read, Int64Array),
DataType::Float32 => limit_array!(array, num_rows_to_read, Float32Array),
DataType::Float64 => limit_array!(array, num_rows_to_read, Float64Array),
DataType::Boolean => limit_array!(array, num_rows_to_read, BooleanArray),
DataType::Utf8 => {
//TODO: this is inefficient and we should improve the Arrow impl to help make this more concise
let b = array.as_any().downcast_ref::<BinaryArray>().unwrap();
let mut values: Vec<String> = Vec::with_capacity(num_rows_to_read as usize);
for i in 0..num_rows_to_read {
values.push(b.get_string(i));
}
let tmp: Vec<&str> = values.iter().map(|s| s.as_str()).collect();
Ok(Arc::new(BinaryArray::from(tmp)))
}
other => Err(ArrowError::ComputeError(format!(
"limit not supported for {:?}",
other
))),
}
}

#[cfg(test)]
mod tests {
use super::*;
Expand Down Expand Up @@ -358,4 +458,37 @@ mod tests {
assert_eq!(5, min(&a).unwrap());
assert_eq!(9, max(&a).unwrap());
}

#[test]
fn test_filter_array() {
let a = Int32Array::from(vec![5, 6, 7, 8, 9]);
let b = BooleanArray::from(vec![true, false, false, true, false]);
let c = filter(&a, &b).unwrap();
let d = c.as_ref().as_any().downcast_ref::<Int32Array>().unwrap();
assert_eq!(2, d.len());
assert_eq!(5, d.value(0));
assert_eq!(8, d.value(1));
}

#[test]
fn test_limit_array() {
let a = Int32Array::from(vec![5, 6, 7, 8, 9]);
let c = limit(&a, 3).unwrap();
let d = c.as_ref().as_any().downcast_ref::<Int32Array>().unwrap();
assert_eq!(3, d.len());
assert_eq!(5, d.value(0));
assert_eq!(6, d.value(1));
assert_eq!(7, d.value(2));
}

#[test]
fn test_limit_array_with_limit_too_large() {
let a = Int32Array::from(vec![5, 6, 7, 8, 9]);
let b = limit(&a, 6);

assert!(
b.is_err(),
"Number of rows to read (6) is larger than the length of the array (5)"
);
}
}
136 changes: 8 additions & 128 deletions rust/datafusion/src/execution/filter.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,8 @@ use std::rc::Rc;
use std::sync::Arc;

use arrow::array::*;
use arrow::datatypes::{DataType, Schema};
use arrow::compute::array_ops::filter;
use arrow::datatypes::Schema;
use arrow::record_batch::RecordBatch;

use super::error::{ExecutionError, Result};
Expand Down Expand Up @@ -61,7 +62,12 @@ impl Relation for FilterRelation {
Some(filter_bools) => {
let filtered_columns: Result<Vec<ArrayRef>> = (0..batch
.num_columns())
.map(|i| filter(batch.column(i), &filter_bools))
.map(|i| {
match filter(batch.column(i).as_ref(), &filter_bools) {
Ok(result) => Ok(result),
Err(error) => Err(ExecutionError::from(error)),
}
})
.collect();

let filtered_batch: RecordBatch = RecordBatch::new(
Expand All @@ -84,129 +90,3 @@ impl Relation for FilterRelation {
&self.schema
}
}

//TODO: move into Arrow array_ops
fn filter(array: &Arc<Array>, filter: &BooleanArray) -> Result<ArrayRef> {
let a = array.as_ref();

//TODO use macros
match a.data_type() {
DataType::UInt8 => {
let b = a.as_any().downcast_ref::<UInt8Array>().unwrap();
let mut builder = UInt8Array::builder(b.len());
for i in 0..b.len() {
if filter.value(i) {
builder.append_value(b.value(i))?;
}
}
Ok(Arc::new(builder.finish()))
}
DataType::UInt16 => {
let b = a.as_any().downcast_ref::<UInt16Array>().unwrap();
let mut builder = UInt16Array::builder(b.len());
for i in 0..b.len() {
if filter.value(i) {
builder.append_value(b.value(i))?;
}
}
Ok(Arc::new(builder.finish()))
}
DataType::UInt32 => {
let b = a.as_any().downcast_ref::<UInt32Array>().unwrap();
let mut builder = UInt32Array::builder(b.len());
for i in 0..b.len() {
if filter.value(i) {
builder.append_value(b.value(i))?;
}
}
Ok(Arc::new(builder.finish()))
}
DataType::UInt64 => {
let b = a.as_any().downcast_ref::<UInt64Array>().unwrap();
let mut builder = UInt64Array::builder(b.len());
for i in 0..b.len() {
if filter.value(i) {
builder.append_value(b.value(i))?;
}
}
Ok(Arc::new(builder.finish()))
}
DataType::Int8 => {
let b = a.as_any().downcast_ref::<Int8Array>().unwrap();
let mut builder = Int8Array::builder(b.len());
for i in 0..b.len() {
if filter.value(i) {
builder.append_value(b.value(i))?;
}
}
Ok(Arc::new(builder.finish()))
}
DataType::Int16 => {
let b = a.as_any().downcast_ref::<Int16Array>().unwrap();
let mut builder = Int16Array::builder(b.len());
for i in 0..b.len() {
if filter.value(i) {
builder.append_value(b.value(i))?;
}
}
Ok(Arc::new(builder.finish()))
}
DataType::Int32 => {
let b = a.as_any().downcast_ref::<Int32Array>().unwrap();
let mut builder = Int32Array::builder(b.len());
for i in 0..b.len() {
if filter.value(i) {
builder.append_value(b.value(i))?;
}
}
Ok(Arc::new(builder.finish()))
}
DataType::Int64 => {
let b = a.as_any().downcast_ref::<Int64Array>().unwrap();
let mut builder = Int64Array::builder(b.len());
for i in 0..b.len() {
if filter.value(i) {
builder.append_value(b.value(i))?;
}
}
Ok(Arc::new(builder.finish()))
}
DataType::Float32 => {
let b = a.as_any().downcast_ref::<Float32Array>().unwrap();
let mut builder = Float32Array::builder(b.len());
for i in 0..b.len() {
if filter.value(i) {
builder.append_value(b.value(i))?;
}
}
Ok(Arc::new(builder.finish()))
}
DataType::Float64 => {
let b = a.as_any().downcast_ref::<Float64Array>().unwrap();
let mut builder = Float64Array::builder(b.len());
for i in 0..b.len() {
if filter.value(i) {
builder.append_value(b.value(i))?;
}
}
Ok(Arc::new(builder.finish()))
}
DataType::Utf8 => {
//TODO: this is inefficient and we should improve the Arrow impl to help make
// this more concise
let b = a.as_any().downcast_ref::<BinaryArray>().unwrap();
let mut values: Vec<String> = Vec::with_capacity(b.len());
for i in 0..b.len() {
if filter.value(i) {
values.push(b.get_string(i));
}
}
let tmp: Vec<&str> = values.iter().map(|s| s.as_str()).collect();
Ok(Arc::new(BinaryArray::from(tmp)))
}
other => Err(ExecutionError::ExecutionError(format!(
"filter not supported for {:?}",
other
))),
}
}
Loading