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

Jump to bottom

ref: Simplify expressions in physical plan #782

Merged
merged 7 commits into from
Oct 2, 2023
Merged

ref: Simplify expressions in physical plan #782

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
cf673d3
ref: Simplify expressions in physical plan
bjchambers Sep 18, 2023
ca4284a
clippy
bjchambers Sep 29, 2023
b974cc1
fix
bjchambers Sep 29, 2023
79b781f
unused imports
bjchambers Sep 29, 2023
5761d80
comments
bjchambers Oct 2, 2023
87d5495
remove input_type
bjchambers Oct 2, 2023
d06ba4c
Merge branch 'main' into ben/updates-to-logical-plans
bjchambers Oct 2, 2023
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
6 changes: 6 additions & 0 deletions Cargo.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -139,7 +139,7 @@ tracing-subscriber = { version = "0.3.17", features = [
"json",
] }
url = { version = "2.3.1", features = ["serde"] }
uuid = { version = "1.3.0", features = ["v4"] }
uuid = { version = "1.3.0", features = ["v4", "serde" ] }

[workspace.dependencies.rocksdb]
# This disables compression algorithms that cause issues during linking due to
Expand Down
133 changes: 51 additions & 82 deletions crates/sparrow-arrow/src/batch.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,10 +1,7 @@
use arrow::error::ArrowError;
use arrow_array::types::{TimestampNanosecondType, UInt64Type};
use arrow_array::{
Array, ArrayRef, ArrowPrimitiveType, BooleanArray, RecordBatch, TimestampNanosecondArray,
UInt64Array,
Array, ArrayRef, ArrowPrimitiveType, BooleanArray, TimestampNanosecondArray, UInt64Array,
};
use arrow_schema::SchemaRef;
use error_stack::{IntoReport, ResultExt};
use itertools::Itertools;

Expand Down Expand Up @@ -36,15 +33,11 @@ impl Batch {

pub fn num_rows(&self) -> usize {
match &self.data {
Some(data) => data.batch.num_rows(),
Some(data) => data.data.len(),
None => 0,
}
}

pub fn into_record_batch(self) -> Option<RecordBatch> {
self.data.map(|data| data.batch)
}

pub fn time(&self) -> Option<&TimestampNanosecondArray> {
self.data.as_ref().map(|data| data.time())
}
Expand All @@ -66,11 +59,8 @@ impl Batch {
self.data.as_ref().map(|info| info.max_present_time)
}

pub fn record_batch(&self) -> Option<&RecordBatch> {
match &self.data {
Some(info) => Some(&info.batch),
None => None,
}
pub fn data(&self) -> Option<&ArrayRef> {
self.data.as_ref().map(|info| &info.data)
}

/// Create a new `Batch` containing the given batch data.
Expand All @@ -83,33 +73,33 @@ impl Batch {
/// (b) all rows in this batch or less than or equal to `up_to_time`
/// (c) all future rows are greater than `up_to_time`.
pub fn new_with_data(
batch: RecordBatch,
data: ArrayRef,
time: ArrayRef,
subsort: ArrayRef,
key_hash: ArrayRef,
up_to_time: RowTime,
) -> Self {
debug_assert_eq!(batch.num_rows(), time.len());
debug_assert_eq!(batch.num_rows(), subsort.len());
debug_assert_eq!(batch.num_rows(), key_hash.len());
debug_assert_eq!(data.len(), time.len());
debug_assert_eq!(data.len(), subsort.len());
debug_assert_eq!(data.len(), key_hash.len());
debug_assert_eq!(time.data_type(), &TimestampNanosecondType::DATA_TYPE);
debug_assert_eq!(subsort.data_type(), &UInt64Type::DATA_TYPE);
debug_assert_eq!(key_hash.data_type(), &UInt64Type::DATA_TYPE);
let data = if batch.num_rows() == 0 {
let data = if data.len() == 0 {
None
} else {
Some(BatchInfo::new(batch, time, subsort, key_hash))
Some(BatchInfo::new(data, time, subsort, key_hash))
};

Self { data, up_to_time }
}

/// Return a new `Batch` with the same time properties, but new data.
pub fn with_projection(&self, new_batch: RecordBatch) -> Self {
assert_eq!(new_batch.num_rows(), self.num_rows());
pub fn with_projection(&self, new_data: ArrayRef) -> Self {
assert_eq!(new_data.len(), self.num_rows());
Self {
data: self.data.as_ref().map(|data| BatchInfo {
batch: new_batch,
data: new_data,
time: data.time.clone(),
subsort: data.subsort.clone(),
key_hash: data.key_hash.clone(),
Expand Down Expand Up @@ -160,11 +150,7 @@ impl Batch {
})
}

pub fn concat(
schema: &SchemaRef,
batches: Vec<Batch>,
up_to_time: RowTime,
) -> error_stack::Result<Batch, Error> {
pub fn concat(batches: Vec<Batch>, up_to_time: RowTime) -> error_stack::Result<Batch, Error> {
// TODO: Add debug assertions for batch ordering?
if batches.iter().all(|batch| batch.is_empty()) {
return Ok(Batch::new_empty(up_to_time));
Expand Down Expand Up @@ -209,16 +195,17 @@ impl Batch {
.change_context(Error::Internal)?;

let batches: Vec<_> = batches
.into_iter()
.flat_map(|batch| batch.into_record_batch())
.iter()
.flat_map(|batch| batch.data())
.map(|data| data.as_ref())
.collect();
let batch = arrow_select::concat::concat_batches(schema, &batches)
let data = arrow_select::concat::concat(&batches)
.into_report()
.change_context(Error::Internal)?;

Ok(Self {
data: Some(BatchInfo {
batch,
data,
time,
subsort,
key_hash,
Expand All @@ -232,8 +219,9 @@ impl Batch {
pub fn take(&self, indices: &UInt64Array) -> error_stack::Result<Self, Error> {
match &self.data {
Some(info) => {
let batch =
take_record_batch(&info.batch, indices).change_context(Error::Internal)?;
let data = arrow_select::take::take(info.data.as_ref(), indices, None)
.into_report()
.change_context(Error::Internal)?;
let time = arrow_select::take::take(info.time.as_ref(), indices, None)
.into_report()
.change_context(Error::Internal)?;
Expand All @@ -244,7 +232,7 @@ impl Batch {
.into_report()
.change_context(Error::Internal)?;
let info = BatchInfo {
batch,
data,
time,
subsort,
key_hash,
Expand Down Expand Up @@ -272,16 +260,8 @@ impl Batch {
let filter = arrow_select::filter::FilterBuilder::new(predicate)
.optimize()
.build();
let columns: Vec<_> = info
.batch
.columns()
.iter()
.map(|column| filter.filter(column))
.try_collect()
.into_report()
.change_context(Error::Internal)?;

let batch = RecordBatch::try_new(info.batch.schema(), columns)
let data = filter
.filter(&info.data)
.into_report()
.change_context(Error::Internal)?;

Expand All @@ -300,7 +280,7 @@ impl Batch {
.into_report()
.change_context(Error::Internal)?;
let info = BatchInfo {
batch,
data,
time,
subsort,
key_hash,
Expand All @@ -325,7 +305,7 @@ impl Batch {
match &self.data {
Some(info) => {
let info = BatchInfo {
batch: info.batch.slice(offset, length),
data: info.data.slice(offset, length),
time: info.time.slice(offset, length),
subsort: info.subsort.slice(offset, length),
key_hash: info.key_hash.slice(offset, length),
Expand Down Expand Up @@ -354,6 +334,8 @@ impl Batch {
) -> Self {
use std::sync::Arc;

use arrow_array::StructArray;

let time: TimestampNanosecondArray = time.into();
let subsort: UInt64Array = (0..(time.len() as u64)).collect_vec().into();
let key_hash: UInt64Array = key_hash.into();
Expand All @@ -362,12 +344,14 @@ impl Batch {
let subsort: ArrayRef = Arc::new(subsort);
let key_hash: ArrayRef = Arc::new(key_hash);

let batch =
RecordBatch::try_new(MINIMAL_SCHEMA.clone(), vec![time.clone(), key_hash.clone()])
.unwrap();
let data = Arc::new(StructArray::new(
MINIMAL_SCHEMA.fields().clone(),
vec![time.clone(), key_hash.clone()],
None,
));

Batch::new_with_data(
batch,
data,
time,
subsort,
key_hash,
Expand All @@ -378,7 +362,7 @@ impl Batch {

#[derive(Clone, Debug)]
pub(crate) struct BatchInfo {
pub(crate) batch: RecordBatch,
pub(crate) data: ArrayRef,
pub(crate) time: ArrayRef,
pub(crate) subsort: ArrayRef,
pub(crate) key_hash: ArrayRef,
Expand All @@ -388,19 +372,19 @@ pub(crate) struct BatchInfo {

impl PartialEq for BatchInfo {
fn eq(&self, other: &Self) -> bool {
self.batch == other.batch
self.data.as_ref() == other.data.as_ref()
&& self.time.as_ref() == other.time.as_ref()
&& self.min_present_time == other.min_present_time
&& self.max_present_time == other.max_present_time
}
}

impl BatchInfo {
fn new(batch: RecordBatch, time: ArrayRef, subsort: ArrayRef, key_hash: ArrayRef) -> Self {
debug_assert_eq!(batch.num_rows(), time.len());
debug_assert_eq!(batch.num_rows(), subsort.len());
debug_assert_eq!(batch.num_rows(), key_hash.len());
debug_assert!(batch.num_rows() > 0);
fn new(data: ArrayRef, time: ArrayRef, subsort: ArrayRef, key_hash: ArrayRef) -> Self {
debug_assert_eq!(data.len(), time.len());
debug_assert_eq!(data.len(), subsort.len());
debug_assert_eq!(data.len(), key_hash.len());
debug_assert!(data.len() > 0);
debug_assert_eq!(time.null_count(), 0);

let time_column: &TimestampNanosecondArray =
Expand All @@ -417,7 +401,7 @@ impl BatchInfo {
RowTime::from_timestamp_ns(time_column.values()[time_column.len() - 1]);

Self {
batch,
data,
time,
subsort,
key_hash,
Expand All @@ -439,21 +423,21 @@ impl BatchInfo {
let slice_start = time_column
.values()
.partition_point(|time| RowTime::from_timestamp_ns(*time) <= time_inclusive);
let slice_len = self.batch.num_rows() - slice_start;
let slice_len = self.data.len() - slice_start;

let max_result_time = RowTime::from_timestamp_ns(time_column.values()[slice_start - 1]);
let min_self_time = RowTime::from_timestamp_ns(time_column.values()[slice_start]);

let result = Self {
batch: self.batch.slice(0, slice_start),
data: self.data.slice(0, slice_start),
time: self.time.slice(0, slice_start),
subsort: self.subsort.slice(0, slice_start),
key_hash: self.key_hash.slice(0, slice_start),
min_present_time: self.min_present_time,
max_present_time: max_result_time,
};

self.batch = self.batch.slice(slice_start, slice_len);
self.data = self.data.slice(slice_start, slice_len);
self.time = self.time.slice(slice_start, slice_len);
self.key_hash = self.key_hash.slice(slice_start, slice_len);
self.min_present_time = min_self_time;
Expand Down Expand Up @@ -482,20 +466,6 @@ pub enum Error {

impl error_stack::Context for Error {}

pub fn take_record_batch(
batch: &RecordBatch,
indices: &UInt64Array,
) -> error_stack::Result<RecordBatch, ArrowError> {
// Produce batches based on indices
let columns = batch
.columns()
.iter()
.map(|c| arrow_select::take::take(c.as_ref(), indices, None).into_report())
.try_collect()?;

RecordBatch::try_new(batch.schema(), columns).into_report()
}

#[cfg(any(test, feature = "testing"))]
#[static_init::dynamic]
static MINIMAL_SCHEMA: arrow_schema::SchemaRef = {
Expand All @@ -517,7 +487,6 @@ mod tests {

use crate::testing::arb_arrays::arb_batch;
use crate::{Batch, RowTime};
use arrow_select::concat::concat_batches;
use itertools::Itertools;
use proptest::prelude::*;

Expand Down Expand Up @@ -546,18 +515,18 @@ mod tests {
prop_assert_eq!(remainder.up_to_time, original.up_to_time);

// create the concatenated result
let concatenated = match (result.record_batch(), remainder.record_batch()) {
let concatenated = match (result.data(), remainder.data()) {
(None, None) => unreachable!(),
(Some(a), None) => a.clone(),
(None, Some(b)) => b.clone(),
(Some(a), Some(b)) => concat_batches(&a.schema(), &[a.clone(), b.clone()]).unwrap(),
(Some(a), Some(b)) => arrow_select::concat::concat(&[a.as_ref(), b.as_ref()]).unwrap(),
};
prop_assert_eq!(&concatenated, original.record_batch().unwrap());
prop_assert_eq!(&concatenated, original.data().unwrap());

prop_assert!(result.data.is_some());
let result = result.data.unwrap();

prop_assert_eq!(result.batch.num_rows(), result.time.len());
prop_assert_eq!(result.data.len(), result.time.len());
prop_assert_eq!(result.time().values()[0], i64::from(result.min_present_time));
prop_assert_eq!(result.time().values()[result.time.len() - 1], i64::from(result.max_present_time));

Expand All @@ -568,7 +537,7 @@ mod tests {
prop_assert!(remainder.data.is_some());
let remainder = remainder.data.unwrap();

prop_assert_eq!(remainder.batch.num_rows(), remainder.time.len());
prop_assert_eq!(remainder.data.len(), remainder.time.len());
prop_assert_eq!(remainder.time().values()[0], i64::from(remainder.min_present_time));
prop_assert_eq!(remainder.time().values()[remainder.time.len() - 1], i64::from(remainder.max_present_time));
}
Expand Down
1 change: 1 addition & 0 deletions crates/sparrow-backend/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ arrow-schema.workspace = true
index_vec.workspace = true
sparrow-core = { path = "../sparrow-core" }
sparrow-physical = { path = "../sparrow-physical" }
uuid.workspace = true

[dev-dependencies]

Expand Down
Loading
Loading