refactor: create a rust based PartitionSet

src/daft-connect/Cargo.toml

@@ -12,7 +12,6 @@ daft-scan = {workspace = true}
 daft-schema = {workspace = true}
 daft-table = {workspace = true}
 dashmap = "6.1.0"
-derive_more = {workspace = true}
 eyre = "0.6.12"
 futures = "0.3.31"
 itertools = {workspace = true}

src/daft-connect/src/translation/logical_plan.rs

@@ -1,8 +1,5 @@
-use std::{collections::HashMap, sync::Arc};
-
 use daft_logical_plan::LogicalPlanBuilder;
-use daft_micropartition::MicroPartition;
-use derive_more::Constructor;
+use daft_micropartition::partitioning::InMemoryPartitionSet;
 use eyre::{bail, Context};
 use spark_connect::{relation::RelType, Limit, Relation};
 use tracing::warn;
@@ -19,17 +16,25 @@
 mod to_df;
 mod with_columns;
 
-#[derive(Constructor)]
 pub struct Plan {
     pub builder: LogicalPlanBuilder,
-    pub psets: HashMap<String, Vec<Arc<MicroPartition>>>,
+    pub psets: InMemoryPartitionSet,
+}
+
+impl Plan {
+    pub fn new(builder: LogicalPlanBuilder) -> Self {
+        Self {
+            builder,
+            psets: InMemoryPartitionSet::default(),
+        }
+    }
 }
 
 impl From<LogicalPlanBuilder> for Plan {
     fn from(builder: LogicalPlanBuilder) -> Self {
         Self {
             builder,
-            psets: HashMap::new(),
+            psets: InMemoryPartitionSet::default(),
         }
     }
 }

src/daft-connect/src/translation/logical_plan/local_relation.rs

@@ -9,6 +9,7 @@ use daft_logical_plan::{
     logical_plan::Source, InMemoryInfo, LogicalPlan, LogicalPlanBuilder, PyLogicalPlanBuilder,
     SourceInfo,
 };
+use daft_micropartition::partitioning::InMemoryPartitionSet;
 use daft_schema::dtype::DaftDataType;
 use daft_table::Table;
 use eyre::{bail, ensure, WrapErr};
@@ -180,7 +181,10 @@ pub fn local_relation(plan: spark_connect::LocalRelation) -> eyre::Result<Plan>
         let mut psets = HashMap::new();
         psets.insert(cache_key, vec![micro_partition]);
 
-        let plan = Plan::new(plan, psets);
+        let plan = Plan {
+            builder: plan,
+            psets: InMemoryPartitionSet::new(psets),
+        };
 
         Ok(plan)
     }

src/daft-local-execution/src/pipeline.rs

@@ -1,4 +1,4 @@
-use std::{collections::HashMap, sync::Arc};
+use std::sync::Arc;
 
 use common_daft_config::DaftExecutionConfig;
 use common_display::{mermaid::MermaidDisplayVisitor, tree::TreeDisplay};
@@ -16,7 +16,7 @@ use daft_local_plan::{
     Sort, UnGroupedAggregate, Unpivot,
 };
 use daft_logical_plan::{stats::StatsState, JoinType};
-use daft_micropartition::MicroPartition;
+use daft_micropartition::{partitioning::PartitionSet, MicroPartition};
 use daft_physical_plan::{extract_agg_expr, populate_aggregation_stages};
 use daft_scan::ScanTaskRef;
 use daft_writers::make_physical_writer_factory;
@@ -75,7 +75,7 @@ pub fn viz_pipeline(root: &dyn PipelineNode) -> String {
 
 pub fn physical_plan_to_pipeline(
     physical_plan: &LocalPhysicalPlan,
-    psets: &HashMap<String, Vec<Arc<MicroPartition>>>,
+    psets: &(impl PartitionSet + ?Sized),
     cfg: &Arc<DaftExecutionConfig>,
 ) -> crate::Result<Box<dyn PipelineNode>> {
     use daft_local_plan::PhysicalScan;
@@ -102,10 +102,10 @@ pub fn physical_plan_to_pipeline(
             scan_task_source.arced().into()
         }
         LocalPhysicalPlan::InMemoryScan(InMemoryScan { info, .. }) => {
-            let partitions = psets
-                .get(&info.cache_key)
-                .unwrap_or_else(|| panic!("Cache key not found: {:?}", info.cache_key));
-            InMemorySource::new(partitions.clone(), info.source_schema.clone())
+            let materialized_pset = psets
+                .get_partition(&info.cache_key)
+                .unwrap_or_else(|_| panic!("Cache key not found: {:?}", info.cache_key));
+            InMemorySource::new(materialized_pset, info.source_schema.clone())
                 .arced()
                 .into()
         }

src/daft-local-execution/src/run.rs

@@ -11,7 +11,10 @@ use common_error::DaftResult;
 use common_tracing::refresh_chrome_trace;
 use daft_local_plan::{translate, LocalPhysicalPlan};
 use daft_logical_plan::LogicalPlanBuilder;
-use daft_micropartition::MicroPartition;
+use daft_micropartition::{
+    partitioning::{InMemoryPartitionSet, PartitionSet},
+    MicroPartition,
+};
 use futures::{FutureExt, Stream};
 use loole::RecvFuture;
 use tokio_util::sync::CancellationToken;
@@ -89,9 +92,10 @@ impl PyNativeExecutor {
                 )
             })
             .collect();
+        let psets = InMemoryPartitionSet::new(native_psets);
         let out = py.allow_threads(|| {
             self.executor
-                .run(native_psets, cfg.config, results_buffer_size)
+                .run(psets, cfg.config, results_buffer_size)
                 .map(|res| res.into_iter())
         })?;
         let iter = Box::new(out.map(|part| {
@@ -121,7 +125,7 @@ impl NativeExecutor {
 
     pub fn run(
         &self,
-        psets: HashMap<String, Vec<Arc<MicroPartition>>>,
+        psets: impl PartitionSet,
         cfg: Arc<DaftExecutionConfig>,
         results_buffer_size: Option<usize>,
     ) -> DaftResult<ExecutionEngineResult> {
@@ -246,7 +250,7 @@ impl IntoIterator for ExecutionEngineResult {
 
 pub fn run_local(
     physical_plan: &LocalPhysicalPlan,
-    psets: HashMap<String, Vec<Arc<MicroPartition>>>,
+    psets: impl PartitionSet,
     cfg: Arc<DaftExecutionConfig>,
     results_buffer_size: Option<usize>,
     cancel: CancellationToken,

src/daft-local-execution/src/sources/in_memory.rs

@@ -4,19 +4,19 @@ use async_trait::async_trait;
 use common_error::DaftResult;
 use daft_core::prelude::SchemaRef;
 use daft_io::IOStatsRef;
-use daft_micropartition::MicroPartition;
+use daft_micropartition::partitioning::PartitionBatchRef;
 use tracing::instrument;
 
 use super::source::Source;
 use crate::sources::source::SourceStream;
 
 pub struct InMemorySource {
-    data: Vec<Arc<MicroPartition>>,
+    data: PartitionBatchRef,
     schema: SchemaRef,
 }
 
 impl InMemorySource {
-    pub fn new(data: Vec<Arc<MicroPartition>>, schema: SchemaRef) -> Self {
+    pub fn new(data: PartitionBatchRef, schema: SchemaRef) -> Self {
         Self { data, schema }
     }
     pub fn arced(self) -> Arc<dyn Source> {
@@ -32,9 +32,7 @@ impl Source for InMemorySource {
         _maintain_order: bool,
         _io_stats: IOStatsRef,
     ) -> DaftResult<SourceStream<'static>> {
-        Ok(Box::pin(futures::stream::iter(
-            self.data.clone().into_iter().map(Ok),
-        )))
+        Ok(self.data.clone().into_partition_stream())
     }
     fn name(&self) -> &'static str {
         "InMemory"

src/daft-micropartition/Cargo.toml

@@ -14,6 +14,7 @@ daft-parquet = {path = "../daft-parquet", default-features = false}
 daft-scan = {path = "../daft-scan", default-features = false}
 daft-stats = {path = "../daft-stats", default-features = false}
 daft-table = {path = "../daft-table", default-features = false}
+futures = {workspace = true}
 parquet2 = {workspace = true}
 pyo3 = {workspace = true, optional = true}
 snafu = {workspace = true}

src/daft-micropartition/src/lib.rs

@@ -16,6 +16,8 @@ use pyo3::PyErr;
 #[cfg(feature = "python")]
 pub use python::register_modules;
 
+pub mod partitioning;
+
 #[derive(Debug, Snafu)]
 pub enum Error {
     #[snafu(display("DaftCoreComputeError: {}", source))]

src/daft-micropartition/src/micropartition.rs

@@ -630,10 +630,12 @@ impl MicroPartition {
             TableState::Loaded(tables) => Ok(tables.clone()),
         }
     }
+
     pub fn get_tables(&self) -> crate::Result<Arc<Vec<Table>>> {
         let tables = self.tables_or_read(IOStatsContext::new("get tables"))?;
         Ok(tables)
     }
+
     pub fn concat_or_get(&self, io_stats: IOStatsRef) -> crate::Result<Arc<Vec<Table>>> {
         let tables = self.tables_or_read(io_stats)?;
         if tables.len() <= 1 {