You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by "tustvold (via GitHub)" <gi...@apache.org> on 2023/04/05 16:57:25 UTC
[GitHub] [arrow-datafusion] tustvold opened a new pull request, #5886: Generify SortPreservingMerge (#5882) (#5879)
tustvold opened a new pull request, #5886:
URL: https://github.com/apache/arrow-datafusion/pull/5886
# Which issue does this PR close?
<!--
We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123.
-->
Part of #5882
Relates to #5879
# Rationale for this change
<!--
Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes.
-->
In order to be able to special case cursors (#5882), we need to first decouple SortPreservingMerge from SortKeyCursor.
# What changes are included in this PR?
<!--
There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR.
-->
* Splits various functionality into smaller, crate-private modules
* Makes SortPreservingMergeStream generic over `Cursor` and accept a type-erased `CursorStream`
* Splits batch construction logic into BatchBuilder
*
# Are these changes tested?
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example, are they covered by existing tests)?
-->
# Are there any user-facing changes?
No
<!--
If there are user-facing changes then we may require documentation to be updated before approving the PR.
-->
<!--
If there are any breaking changes to public APIs, please add the `api change` label.
-->
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] tustvold commented on a diff in pull request #5886: Generify SortPreservingMerge (#5882) (#5879)
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #5886:
URL: https://github.com/apache/arrow-datafusion/pull/5886#discussion_r1159101163
##########
datafusion/core/src/physical_plan/sorts/stream.rs:
##########
@@ -0,0 +1,145 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::common::Result;
+use crate::physical_plan::sorts::cursor::SortKeyCursor;
+use crate::physical_plan::SendableRecordBatchStream;
+use crate::physical_plan::{PhysicalExpr, PhysicalSortExpr};
+use arrow::datatypes::Schema;
+use arrow::record_batch::RecordBatch;
+use arrow::row::{RowConverter, SortField};
+use futures::stream::{Fuse, StreamExt};
+use std::sync::Arc;
+use std::task::{ready, Context, Poll};
+
+/// A [`Stream`](futures::Stream) that has multiple partitions that can
+/// be polled separately but not concurrently
+///
+/// Used by sort preserving merge to decouple the cursor merging logic from
+/// the source of the cursors, the intention being to allow preserving
Review Comment:
Among potentially other things
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] tustvold commented on a diff in pull request #5886: Generify SortPreservingMerge (#5882) (#5879)
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #5886:
URL: https://github.com/apache/arrow-datafusion/pull/5886#discussion_r1158786883
##########
datafusion/core/src/physical_plan/sorts/mod.rs:
##########
@@ -17,30 +17,14 @@
//! Sort functionalities
-use crate::physical_plan::SendableRecordBatchStream;
-use std::fmt::{Debug, Formatter};
-
+mod builder;
mod cursor;
mod index;
+mod merge;
pub mod sort;
pub mod sort_preserving_merge;
+mod stream;
pub use cursor::SortKeyCursor;
pub use index::RowIndex;
-
-pub(crate) struct SortedStream {
Review Comment:
This abstraction was a tad pointless, it was only used by `ExternalSorter` to propagate the size of the in-memory sorted batches. I adjusted it to just call init_mem_used itself, and created #5885 to track the broader pre-existing issue of memory accounting within merge streams
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] tustvold commented on a diff in pull request #5886: Generify SortPreservingMerge (#5882) (#5879)
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #5886:
URL: https://github.com/apache/arrow-datafusion/pull/5886#discussion_r1159101027
##########
datafusion/core/src/physical_plan/sorts/cursor.rs:
##########
@@ -110,3 +110,24 @@ impl Ord for SortKeyCursor {
}
}
}
+
Review Comment:
There will be implementations specialized for the various single column types
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] tustvold commented on pull request #5886: Generify SortPreservingMerge (#5882) (#5879)
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on PR #5886:
URL: https://github.com/apache/arrow-datafusion/pull/5886#issuecomment-1499113385
https://github.com/apache/arrow-datafusion/pull/5894 and https://github.com/apache/arrow-datafusion/pull/5895 both build off this and improve performance by about ~10% each. I will leave this open for a bit longer, but unless anybody objects I intend to merge this tomorrow morning to keep things moving along.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] alamb commented on a diff in pull request #5886: Generify SortPreservingMerge (#5882) (#5879)
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #5886:
URL: https://github.com/apache/arrow-datafusion/pull/5886#discussion_r1159027422
##########
datafusion/core/src/physical_plan/sorts/builder.rs:
##########
@@ -0,0 +1,171 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::common::Result;
+use crate::physical_plan::sorts::index::RowIndex;
+use arrow::array::{make_array, MutableArrayData};
+use arrow::datatypes::SchemaRef;
+use arrow::record_batch::RecordBatch;
+use std::collections::VecDeque;
+
+/// Provides an API to incrementally build a [`RecordBatch`] from partitioned [`RecordBatch`]
Review Comment:
```suggestion
/// incrementally builds output [`RecordBatch`]es from multiple streams of [`RecordBatch`] and
/// using [`RowIndex`] to describe the outputs
```
##########
datafusion/core/src/physical_plan/sorts/builder.rs:
##########
@@ -0,0 +1,171 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::common::Result;
+use crate::physical_plan::sorts::index::RowIndex;
+use arrow::array::{make_array, MutableArrayData};
+use arrow::datatypes::SchemaRef;
+use arrow::record_batch::RecordBatch;
+use std::collections::VecDeque;
+
+/// Provides an API to incrementally build a [`RecordBatch`] from partitioned [`RecordBatch`]
+#[derive(Debug)]
+pub struct BatchBuilder {
+ /// The schema of the RecordBatches yielded by this stream
+ schema: SchemaRef,
+ /// For each input stream maintain a dequeue of RecordBatches
+ ///
+ /// Exhausted batches will be popped off the front once all
+ /// their rows have been yielded to the output
+ batches: Vec<VecDeque<RecordBatch>>,
+
+ /// The accumulated row indexes for the next record batch
+ indices: Vec<RowIndex>,
+}
+
+impl BatchBuilder {
+ /// Create a new [`BatchBuilder`] with the provided `stream_count` and `batch_size`
+ pub fn new(schema: SchemaRef, stream_count: usize, batch_size: usize) -> Self {
+ let batches = (0..stream_count).map(|_| VecDeque::new()).collect();
+
+ Self {
+ schema,
+ batches,
+ indices: Vec::with_capacity(batch_size),
+ }
+ }
+
+ /// Append a new batch in `stream_idx`
+ pub fn push_batch(&mut self, stream_idx: usize, batch: RecordBatch) {
+ self.batches[stream_idx].push_back(batch)
+ }
+
+ /// Push `row_idx` from the most recently appended batch in `stream_idx`
+ pub fn push_row(&mut self, stream_idx: usize, row_idx: usize) {
+ let batch_idx = self.batches[stream_idx].len() - 1;
+ self.indices.push(RowIndex {
+ stream_idx,
+ batch_idx,
+ row_idx,
+ });
+ }
+
+ /// Returns the number of rows in this [`BatchBuilder`]
Review Comment:
```suggestion
/// Returns the number of in-progress rows in this [`BatchBuilder`].
///
/// This is the maximum number or rows that can be created by calling [`build_record_batch`]
```
##########
datafusion/core/src/physical_plan/sorts/stream.rs:
##########
@@ -0,0 +1,145 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::common::Result;
+use crate::physical_plan::sorts::cursor::SortKeyCursor;
+use crate::physical_plan::SendableRecordBatchStream;
+use crate::physical_plan::{PhysicalExpr, PhysicalSortExpr};
+use arrow::datatypes::Schema;
+use arrow::record_batch::RecordBatch;
+use arrow::row::{RowConverter, SortField};
+use futures::stream::{Fuse, StreamExt};
+use std::sync::Arc;
+use std::task::{ready, Context, Poll};
+
+/// A [`Stream`](futures::Stream) that has multiple partitions that can
+/// be polled separately but not concurrently
+///
+/// Used by sort preserving merge to decouple the cursor merging logic from
+/// the source of the cursors, the intention being to allow preserving
+/// any row encoding performed for intermediate sorts
+pub trait PartitionedStream: std::fmt::Debug + Send {
+ type Output;
+
+ /// Returns the number of partitions
+ fn partitions(&self) -> usize;
+
+ fn poll_next(
Review Comment:
Is there a reason this implements `poll_next` itself (rather than extending the standard `Stream` trait directly -- like `: Stream`?
##########
datafusion/core/src/physical_plan/sorts/builder.rs:
##########
@@ -0,0 +1,171 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::common::Result;
+use crate::physical_plan::sorts::index::RowIndex;
+use arrow::array::{make_array, MutableArrayData};
+use arrow::datatypes::SchemaRef;
+use arrow::record_batch::RecordBatch;
+use std::collections::VecDeque;
+
+/// Provides an API to incrementally build a [`RecordBatch`] from partitioned [`RecordBatch`]
+#[derive(Debug)]
+pub struct BatchBuilder {
+ /// The schema of the RecordBatches yielded by this stream
+ schema: SchemaRef,
+ /// For each input stream maintain a dequeue of RecordBatches
+ ///
+ /// Exhausted batches will be popped off the front once all
+ /// their rows have been yielded to the output
+ batches: Vec<VecDeque<RecordBatch>>,
+
+ /// The accumulated row indexes for the next record batch
+ indices: Vec<RowIndex>,
+}
+
+impl BatchBuilder {
+ /// Create a new [`BatchBuilder`] with the provided `stream_count` and `batch_size`
+ pub fn new(schema: SchemaRef, stream_count: usize, batch_size: usize) -> Self {
+ let batches = (0..stream_count).map(|_| VecDeque::new()).collect();
+
+ Self {
+ schema,
+ batches,
+ indices: Vec::with_capacity(batch_size),
+ }
+ }
+
+ /// Append a new batch in `stream_idx`
+ pub fn push_batch(&mut self, stream_idx: usize, batch: RecordBatch) {
+ self.batches[stream_idx].push_back(batch)
+ }
+
+ /// Push `row_idx` from the most recently appended batch in `stream_idx`
Review Comment:
Also, I think it is a requirement that the row_idx are pushed in order, right?
##########
datafusion/core/src/physical_plan/sorts/cursor.rs:
##########
@@ -110,3 +110,24 @@ impl Ord for SortKeyCursor {
}
}
}
+
Review Comment:
What is the plan for the `Cursor` trait? As in why make at trait that has only a single impl?
##########
datafusion/core/src/physical_plan/sorts/stream.rs:
##########
@@ -0,0 +1,145 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::common::Result;
+use crate::physical_plan::sorts::cursor::SortKeyCursor;
+use crate::physical_plan::SendableRecordBatchStream;
+use crate::physical_plan::{PhysicalExpr, PhysicalSortExpr};
+use arrow::datatypes::Schema;
+use arrow::record_batch::RecordBatch;
+use arrow::row::{RowConverter, SortField};
+use futures::stream::{Fuse, StreamExt};
+use std::sync::Arc;
+use std::task::{ready, Context, Poll};
+
+/// A [`Stream`](futures::Stream) that has multiple partitions that can
+/// be polled separately but not concurrently
+///
+/// Used by sort preserving merge to decouple the cursor merging logic from
+/// the source of the cursors, the intention being to allow preserving
Review Comment:
👍
So the idea is there will be a specialized implementation of cursors for row and single column primtives, etc?
##########
datafusion/core/src/physical_plan/sorts/builder.rs:
##########
@@ -0,0 +1,171 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::common::Result;
+use crate::physical_plan::sorts::index::RowIndex;
+use arrow::array::{make_array, MutableArrayData};
+use arrow::datatypes::SchemaRef;
+use arrow::record_batch::RecordBatch;
+use std::collections::VecDeque;
+
+/// Provides an API to incrementally build a [`RecordBatch`] from partitioned [`RecordBatch`]
+#[derive(Debug)]
+pub struct BatchBuilder {
+ /// The schema of the RecordBatches yielded by this stream
+ schema: SchemaRef,
+ /// For each input stream maintain a dequeue of RecordBatches
+ ///
+ /// Exhausted batches will be popped off the front once all
+ /// their rows have been yielded to the output
+ batches: Vec<VecDeque<RecordBatch>>,
+
+ /// The accumulated row indexes for the next record batch
+ indices: Vec<RowIndex>,
+}
+
+impl BatchBuilder {
+ /// Create a new [`BatchBuilder`] with the provided `stream_count` and `batch_size`
+ pub fn new(schema: SchemaRef, stream_count: usize, batch_size: usize) -> Self {
+ let batches = (0..stream_count).map(|_| VecDeque::new()).collect();
+
+ Self {
+ schema,
+ batches,
+ indices: Vec::with_capacity(batch_size),
+ }
+ }
+
+ /// Append a new batch in `stream_idx`
+ pub fn push_batch(&mut self, stream_idx: usize, batch: RecordBatch) {
+ self.batches[stream_idx].push_back(batch)
+ }
+
+ /// Push `row_idx` from the most recently appended batch in `stream_idx`
Review Comment:
```suggestion
/// Push `row_idx` from the most recently appended batch in `stream_idx`
/// into the queue of rows to generate as output
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] tustvold commented on a diff in pull request #5886: Generify SortPreservingMerge (#5882) (#5879)
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #5886:
URL: https://github.com/apache/arrow-datafusion/pull/5886#discussion_r1159100763
##########
datafusion/core/src/physical_plan/sorts/builder.rs:
##########
@@ -0,0 +1,171 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::common::Result;
+use crate::physical_plan::sorts::index::RowIndex;
+use arrow::array::{make_array, MutableArrayData};
+use arrow::datatypes::SchemaRef;
+use arrow::record_batch::RecordBatch;
+use std::collections::VecDeque;
+
+/// Provides an API to incrementally build a [`RecordBatch`] from partitioned [`RecordBatch`]
+#[derive(Debug)]
+pub struct BatchBuilder {
+ /// The schema of the RecordBatches yielded by this stream
+ schema: SchemaRef,
+ /// For each input stream maintain a dequeue of RecordBatches
+ ///
+ /// Exhausted batches will be popped off the front once all
+ /// their rows have been yielded to the output
+ batches: Vec<VecDeque<RecordBatch>>,
+
+ /// The accumulated row indexes for the next record batch
+ indices: Vec<RowIndex>,
+}
+
+impl BatchBuilder {
+ /// Create a new [`BatchBuilder`] with the provided `stream_count` and `batch_size`
+ pub fn new(schema: SchemaRef, stream_count: usize, batch_size: usize) -> Self {
+ let batches = (0..stream_count).map(|_| VecDeque::new()).collect();
+
+ Self {
+ schema,
+ batches,
+ indices: Vec::with_capacity(batch_size),
+ }
+ }
+
+ /// Append a new batch in `stream_idx`
+ pub fn push_batch(&mut self, stream_idx: usize, batch: RecordBatch) {
+ self.batches[stream_idx].push_back(batch)
+ }
+
+ /// Push `row_idx` from the most recently appended batch in `stream_idx`
Review Comment:
Actually the way this is implemented does not require this
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] tustvold commented on a diff in pull request #5886: Generify SortPreservingMerge (#5882) (#5879)
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #5886:
URL: https://github.com/apache/arrow-datafusion/pull/5886#discussion_r1159101332
##########
datafusion/core/src/physical_plan/sorts/stream.rs:
##########
@@ -0,0 +1,145 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::common::Result;
+use crate::physical_plan::sorts::cursor::SortKeyCursor;
+use crate::physical_plan::SendableRecordBatchStream;
+use crate::physical_plan::{PhysicalExpr, PhysicalSortExpr};
+use arrow::datatypes::Schema;
+use arrow::record_batch::RecordBatch;
+use arrow::row::{RowConverter, SortField};
+use futures::stream::{Fuse, StreamExt};
+use std::sync::Arc;
+use std::task::{ready, Context, Poll};
+
+/// A [`Stream`](futures::Stream) that has multiple partitions that can
+/// be polled separately but not concurrently
+///
+/// Used by sort preserving merge to decouple the cursor merging logic from
+/// the source of the cursors, the intention being to allow preserving
+/// any row encoding performed for intermediate sorts
+pub trait PartitionedStream: std::fmt::Debug + Send {
+ type Output;
+
+ /// Returns the number of partitions
+ fn partitions(&self) -> usize;
+
+ fn poll_next(
Review Comment:
Because it takes a stream index
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] tustvold commented on pull request #5886: Generify SortPreservingMerge (#5882) (#5879)
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on PR #5886:
URL: https://github.com/apache/arrow-datafusion/pull/5886#issuecomment-1499365858
https://github.com/apache/arrow-datafusion/pull/5897/commits/5f7a3d6bdd057865207cee35b3e7f81071ae16bb#diff-95c601f2909c85052924354d8d161a9f7bf539c47b59abb1d525a6d217b4e402R60-R61 shows how this can be used to specialize merge for a single primitive column
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] tustvold commented on a diff in pull request #5886: Generify SortPreservingMerge (#5882) (#5879)
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #5886:
URL: https://github.com/apache/arrow-datafusion/pull/5886#discussion_r1158785515
##########
datafusion/core/src/physical_plan/sorts/merge.rs:
##########
@@ -0,0 +1,258 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::common::Result;
+use crate::physical_plan::metrics::MemTrackingMetrics;
+use crate::physical_plan::sorts::builder::BatchBuilder;
+use crate::physical_plan::sorts::cursor::Cursor;
+use crate::physical_plan::sorts::stream::{PartitionedStream, SortKeyCursorStream};
+use crate::physical_plan::{
+ PhysicalSortExpr, RecordBatchStream, SendableRecordBatchStream,
+};
+use arrow::datatypes::SchemaRef;
+use arrow::record_batch::RecordBatch;
+use futures::Stream;
+use std::pin::Pin;
+use std::task::{ready, Context, Poll};
+
+/// Perform a streaming merge of [`SendableRecordBatchStream`]
+pub(crate) fn streaming_merge(
+ streams: Vec<SendableRecordBatchStream>,
+ schema: SchemaRef,
+ expressions: &[PhysicalSortExpr],
+ tracking_metrics: MemTrackingMetrics,
+ batch_size: usize,
+) -> Result<SendableRecordBatchStream> {
+ let streams = SortKeyCursorStream::try_new(schema.as_ref(), expressions, streams)?;
+
+ Ok(Box::pin(SortPreservingMergeStream::new(
+ Box::new(streams),
+ schema,
+ tracking_metrics,
+ batch_size,
+ )))
+}
+
+/// A fallible [`PartitionedStream`] of [`Cursor`] and [`RecordBatch`]
+type CursorStream<C> = Box<dyn PartitionedStream<Output = Result<(C, RecordBatch)>>>;
+
+#[derive(Debug)]
+struct SortPreservingMergeStream<C> {
Review Comment:
This is ported from sort_preserving_merge.rs but tweaked heavily to make it slightly easier to follow (hopefully)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] tustvold commented on a diff in pull request #5886: Generify SortPreservingMerge (#5882) (#5879)
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #5886:
URL: https://github.com/apache/arrow-datafusion/pull/5886#discussion_r1158784705
##########
datafusion/core/src/physical_plan/sorts/builder.rs:
##########
@@ -0,0 +1,171 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::common::Result;
+use crate::physical_plan::sorts::index::RowIndex;
+use arrow::array::{make_array, MutableArrayData};
+use arrow::datatypes::SchemaRef;
+use arrow::record_batch::RecordBatch;
+use std::collections::VecDeque;
+
+/// Provides an API to incrementally build a [`RecordBatch`] from partitioned [`RecordBatch`]
+#[derive(Debug)]
+pub struct BatchBuilder {
+ /// The schema of the RecordBatches yielded by this stream
+ schema: SchemaRef,
+ /// For each input stream maintain a dequeue of RecordBatches
+ ///
+ /// Exhausted batches will be popped off the front once all
+ /// their rows have been yielded to the output
+ batches: Vec<VecDeque<RecordBatch>>,
+
+ /// The accumulated row indexes for the next record batch
+ indices: Vec<RowIndex>,
+}
+
+impl BatchBuilder {
+ /// Create a new [`BatchBuilder`] with the provided `stream_count` and `batch_size`
+ pub fn new(schema: SchemaRef, stream_count: usize, batch_size: usize) -> Self {
+ let batches = (0..stream_count).map(|_| VecDeque::new()).collect();
+
+ Self {
+ schema,
+ batches,
+ indices: Vec::with_capacity(batch_size),
+ }
+ }
+
+ /// Append a new batch in `stream_idx`
+ pub fn push_batch(&mut self, stream_idx: usize, batch: RecordBatch) {
+ self.batches[stream_idx].push_back(batch)
+ }
+
+ /// Push `row_idx` from the most recently appended batch in `stream_idx`
+ pub fn push_row(&mut self, stream_idx: usize, row_idx: usize) {
+ let batch_idx = self.batches[stream_idx].len() - 1;
+ self.indices.push(RowIndex {
+ stream_idx,
+ batch_idx,
+ row_idx,
+ });
+ }
+
+ /// Returns the number of rows in this [`BatchBuilder`]
+ pub fn len(&self) -> usize {
+ self.indices.len()
+ }
+
+ /// Returns `true` if this [`BatchBuilder`] contains no in-progress rows
+ pub fn is_empty(&self) -> bool {
+ self.indices.is_empty()
+ }
+
+ /// Returns the schema of this [`BatchBuilder`]
+ pub fn schema(&self) -> &SchemaRef {
+ &self.schema
+ }
+
+ /// Drains the in_progress row indexes, and builds a new RecordBatch from them
+ ///
+ /// Will then drop any batches for which all rows have been yielded to the output
+ ///
+ /// Returns `None` if no pending rows
+ pub fn build_record_batch(&mut self) -> Result<Option<RecordBatch>> {
Review Comment:
This method is moved pretty much verbatim from sort_preserving_merge.rs
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] tustvold commented on pull request #5886: Generify SortPreservingMerge (#5882) (#5879)
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on PR #5886:
URL: https://github.com/apache/arrow-datafusion/pull/5886#issuecomment-1497834241
I have confirmed this has no discernible impact on the existing benchmarks, nor the benchmarks added in #5881
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] tustvold commented on a diff in pull request #5886: Generify SortPreservingMerge (#5882) (#5879)
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #5886:
URL: https://github.com/apache/arrow-datafusion/pull/5886#discussion_r1159712464
##########
datafusion/core/src/physical_plan/sorts/builder.rs:
##########
@@ -0,0 +1,171 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::common::Result;
+use crate::physical_plan::sorts::index::RowIndex;
+use arrow::array::{make_array, MutableArrayData};
+use arrow::datatypes::SchemaRef;
+use arrow::record_batch::RecordBatch;
+use std::collections::VecDeque;
+
+/// Provides an API to incrementally build a [`RecordBatch`] from partitioned [`RecordBatch`]
Review Comment:
I actually intend for a future PR to remove the use of RowIndex
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] tustvold merged pull request #5886: Generify SortPreservingMerge (#5882) (#5879)
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold merged PR #5886:
URL: https://github.com/apache/arrow-datafusion/pull/5886
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] alamb commented on pull request #5886: Generify SortPreservingMerge (#5882) (#5879)
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #5886:
URL: https://github.com/apache/arrow-datafusion/pull/5886#issuecomment-1498160118
cc @yjshen and @jaylmiller
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] tustvold commented on a diff in pull request #5886: Generify SortPreservingMerge (#5882) (#5879)
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #5886:
URL: https://github.com/apache/arrow-datafusion/pull/5886#discussion_r1159101332
##########
datafusion/core/src/physical_plan/sorts/stream.rs:
##########
@@ -0,0 +1,145 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::common::Result;
+use crate::physical_plan::sorts::cursor::SortKeyCursor;
+use crate::physical_plan::SendableRecordBatchStream;
+use crate::physical_plan::{PhysicalExpr, PhysicalSortExpr};
+use arrow::datatypes::Schema;
+use arrow::record_batch::RecordBatch;
+use arrow::row::{RowConverter, SortField};
+use futures::stream::{Fuse, StreamExt};
+use std::sync::Arc;
+use std::task::{ready, Context, Poll};
+
+/// A [`Stream`](futures::Stream) that has multiple partitions that can
+/// be polled separately but not concurrently
+///
+/// Used by sort preserving merge to decouple the cursor merging logic from
+/// the source of the cursors, the intention being to allow preserving
+/// any row encoding performed for intermediate sorts
+pub trait PartitionedStream: std::fmt::Debug + Send {
+ type Output;
+
+ /// Returns the number of partitions
+ fn partitions(&self) -> usize;
+
+ fn poll_next(
Review Comment:
Because it takes a stream index, we can't just use a `Vec<Stream>` as we need shared mutable state across the streams
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org