Make MemoryManager and MemoryStream public (#1664)

* Make MemoryManager and MemoryStream public * lint
apache · Jan 24, 2022 · 97f95b3 · 97f95b3
1 parent c63cfd4
commit 97f95b3
Show file tree

Hide file tree

Showing 9 changed files with 3 additions and 100 deletions.
diff --git a/ballista/rust/core/src/client.rs b/ballista/rust/core/src/client.rs
@@ -25,7 +25,6 @@ use std::{
 };
 
 use crate::error::{ballista_error, BallistaError, Result};
-use crate::memory_stream::MemoryStream;
 use crate::serde::protobuf::{self};
 use crate::serde::scheduler::{
     Action, ExecutePartition, ExecutePartitionResult, PartitionId, PartitionStats,

diff --git a/ballista/rust/core/src/execution_plans/shuffle_reader.rs b/ballista/rust/core/src/execution_plans/shuffle_reader.rs
@@ -20,7 +20,6 @@ use std::sync::Arc;
 use std::{any::Any, pin::Pin};
 
 use crate::client::BallistaClient;
-use crate::memory_stream::MemoryStream;
 use crate::serde::scheduler::{PartitionLocation, PartitionStats};
 
 use crate::utils::WrappedStream;

diff --git a/ballista/rust/core/src/execution_plans/shuffle_writer.rs b/ballista/rust/core/src/execution_plans/shuffle_writer.rs
@@ -28,7 +28,6 @@ use std::time::Instant;
 use std::{any::Any, pin::Pin};
 
 use crate::error::BallistaError;
-use crate::memory_stream::MemoryStream;
 use crate::utils;
 
 use crate::serde::protobuf::ShuffleWritePartition;
@@ -47,6 +46,7 @@ use datafusion::error::{DataFusionError, Result};
 use datafusion::execution::runtime_env::RuntimeEnv;
 use datafusion::physical_plan::common::IPCWriter;
 use datafusion::physical_plan::hash_utils::create_hashes;
+use datafusion::physical_plan::memory::MemoryStream;
 use datafusion::physical_plan::metrics::{
     self, ExecutionPlanMetricsSet, MetricBuilder, MetricsSet,
 };

diff --git a/ballista/rust/core/src/execution_plans/unresolved_shuffle.rs b/ballista/rust/core/src/execution_plans/unresolved_shuffle.rs
@@ -18,7 +18,6 @@
 use std::sync::Arc;
 use std::{any::Any, pin::Pin};
 
-use crate::memory_stream::MemoryStream;
 use crate::serde::scheduler::PartitionLocation;
 
 use async_trait::async_trait;

diff --git a/ballista/rust/core/src/lib.rs b/ballista/rust/core/src/lib.rs
@@ -27,7 +27,6 @@ pub mod client;
 pub mod config;
 pub mod error;
 pub mod execution_plans;
-pub mod memory_stream;
 pub mod utils;
 
 #[macro_use]

diff --git a/ballista/rust/core/src/memory_stream.rs b/ballista/rust/core/src/memory_stream.rs
diff --git a/ballista/rust/core/src/utils.rs b/ballista/rust/core/src/utils.rs
@@ -26,7 +26,6 @@ use crate::error::{BallistaError, Result};
 use crate::execution_plans::{
     DistributedQueryExec, ShuffleWriterExec, UnresolvedShuffleExec,
 };
-use crate::memory_stream::MemoryStream;
 use crate::serde::scheduler::PartitionStats;
 
 use crate::config::BallistaConfig;

diff --git a/datafusion/src/execution/mod.rs b/datafusion/src/execution/mod.rs
@@ -20,6 +20,6 @@
 pub mod context;
 pub mod dataframe_impl;
 pub(crate) mod disk_manager;
-pub(crate) mod memory_manager;
+pub mod memory_manager;
 pub mod options;
 pub mod runtime_env;
diff --git a/datafusion/src/physical_plan/memory.rs b/datafusion/src/physical_plan/memory.rs
@@ -147,7 +147,7 @@ impl MemoryExec {
 }
 
 /// Iterator over batches
-pub(crate) struct MemoryStream {
+pub struct MemoryStream {
     /// Vector of record batches
     data: Vec<RecordBatch>,
     /// Schema representing the data