feat(connect): df.show

Cargo.toml

@@ -3,6 +3,7 @@ common-daft-config = {path = "src/common/daft-config", default-features = false}
 common-display = {path = "src/common/display", default-features = false}
 common-file-formats = {path = "src/common/file-formats", default-features = false}
 common-hashable-float-wrapper = {path = "src/common/hashable-float-wrapper", default-features = false}
+common-partitioning = {path = "src/common/partitioning", default-features = false}
 common-resource-request = {path = "src/common/resource-request", default-features = false}
 common-runtime = {path = "src/common/runtime", default-features = false}
 common-scan-info = {path = "src/common/scan-info", default-features = false}
@@ -46,18 +47,12 @@ sysinfo = {workspace = true}
 # maturin will turn this on
 python = [
   "common-daft-config/python",
-  "common-daft-config/python",
-  "common-daft-config/python",
-  "common-display/python",
   "common-display/python",
-  "common-display/python",
-  "common-resource-request/python",
-  "common-resource-request/python",
+  "common-partitioning/python",
   "common-resource-request/python",
+  "common-file-formats/python",
   "common-scan-info/python",
   "common-system-info/python",
-  "common-system-info/python",
-  "common-system-info/python",
   "daft-catalog-python-catalog/python",
   "daft-catalog/python",
   "daft-connect/python",
@@ -79,7 +74,6 @@ python = [
   "daft-scheduler/python",
   "daft-sql/python",
   "daft-stats/python",
-  "daft-stats/python",
   "daft-table/python",
   "daft-writers/python",
   "dep:daft-catalog-python-catalog",
@@ -175,7 +169,8 @@ members = [
   "src/daft-connect",
   "src/parquet2",
   # "src/spark-connect-script",
-  "src/generated/spark-connect"
+  "src/generated/spark-connect",
+  "src/common/partitioning"
 ]
 
 [workspace.dependencies]

src/common/partitioning/Cargo.toml

@@ -0,0 +1,17 @@
+[dependencies]
+common-py-serde = {path = "../py-serde", optional = true}
+pyo3 = {workspace = true, optional = true}
+common-error.workspace = true
+futures.workspace = true
+serde.workspace = true
+
+[features]
+python = ["dep:pyo3", "common-error/python", "common-py-serde"]
+
+[lints]
+workspace = true
+
+[package]
+name = "common-partitioning"
+edition.workspace = true
+version.workspace = true

src/common/partitioning/src/lib.rs

@@ -0,0 +1,90 @@
+use std::{any::Any, sync::Arc};
+
+use common_error::DaftResult;
+use futures::stream::BoxStream;
+use serde::{Deserialize, Serialize};
+#[cfg(feature = "python")]
+use {
+    common_py_serde::{deserialize_py_object, serialize_py_object},
+    pyo3::PyObject,
+};
+
+/// Common trait interface for dataset partitioning, defined in this shared crate to avoid circular dependencies.
+/// Acts as a forward reference for concrete partition implementations. _(Specifically the `MicroPartition` type defined in `daft-micropartition`)_
+pub trait Partition: std::fmt::Debug + Send + Sync {
+    fn as_any(&self) -> &dyn std::any::Any;
+    fn as_any_arc(self: Arc<Self>) -> Arc<dyn Any + Send + Sync>;
+    fn size_bytes(&self) -> DaftResult<Option<usize>>;
+}
+
+/// An Arc'd reference to a [`Partition`]
+pub type PartitionRef = Arc<dyn Partition>;
+
+/// Key used to identify a partition
+pub type PartitionId = Arc<str>;
+
+/// A collection of related [`Partitions`]'s that can be processed as a single unit.
+/// All partitions in a batch should have the same schema.
+pub trait PartitionBatch<T: Partition>: std::fmt::Debug + Send + Sync {
+    fn partitions(&self) -> Vec<Arc<T>>;
+    fn metadata(&self) -> PartitionMetadata;
+    /// consume the partition batch and return a stream of partitions
+    fn into_partition_stream(self: Arc<Self>) -> BoxStream<'static, DaftResult<Arc<T>>>;
+}
+
+/// An Arc'd reference to a [`PartitionBatch`]
+pub type PartitionBatchRef<T> = Arc<dyn PartitionBatch<T>>;
+
+/// ported over from `daft/runners/partitioning.py`
+// TODO: port over the rest of the functionality
+#[derive(Debug, Clone)]
+pub struct PartitionMetadata {
+    pub num_rows: usize,
+    pub size_bytes: usize,
+}
+
+/// A partition set is a collection of partition batches.
+/// It is up to the implementation to decide how to store and manage the partition batches.
+/// For example, an in memory partition set could likely be stored as `HashMap<PartitionId, PartitionBatchRef<T>>`.
+///
+/// It is important to note that the methods do not take `&mut self` but instead take `&self`.
+/// So it is up to the implementation to manage any interior mutability.
+pub trait PartitionSet<T: Partition>: std::fmt::Debug + Send + Sync {
+    fn as_any(&self) -> &dyn Any;
+    fn as_any_arc(self: Arc<Self>) -> Arc<dyn Any + Send + Sync>;
+    /// Merge all micropartitions into a single micropartition
+    fn get_merged_partitions(&self) -> DaftResult<PartitionRef>;
+    /// Get a preview of the micropartitions
+    fn get_preview_partitions(&self, num_rows: usize) -> DaftResult<PartitionBatchRef<T>>;
+    /// Number of partitions
+    fn num_partitions(&self) -> usize;
+    fn len(&self) -> usize;
+    /// Check if the partition set is empty
+    fn is_empty(&self) -> bool;
+    /// Size of the partition set in bytes
+    fn size_bytes(&self) -> DaftResult<usize>;
+    /// Check if a partition exists
+    fn has_partition(&self, idx: &PartitionId) -> bool;
+    /// Delete a partition
+    fn delete_partition(&self, idx: &PartitionId) -> DaftResult<()>;
+    /// Set a partition
+    fn set_partition(&self, idx: PartitionId, part: &dyn PartitionBatch<T>) -> DaftResult<()>;
+    /// Get a partition
+    fn get_partition(&self, idx: &PartitionId) -> DaftResult<PartitionBatchRef<T>>;
+
+    /// Consume the partition set and return a stream of partitions
+    fn into_partition_stream(self: Arc<Self>) -> BoxStream<'static, DaftResult<Arc<T>>>;
+}
+
+pub type PartitionSetRef<T> = Arc<dyn PartitionSet<T>>;
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub enum PartitionCacheEntry {
+    #[serde(
+        serialize_with = "serialize_py_object",
+        deserialize_with = "deserialize_py_object"
+    )]
+    #[cfg(feature = "python")]
+    Python(PyObject),
+    Rust(String),
+}

src/daft-connect/Cargo.toml

@@ -3,6 +3,7 @@ arrow2 = {workspace = true, features = ["io_json_integration"]}
 async-stream = "0.3.6"
 color-eyre = "0.6.3"
 common-daft-config = {workspace = true}
+common-error = {workspace = true}
 common-file-formats = {workspace = true}
 daft-core = {workspace = true}
 daft-dsl = {workspace = true}

src/daft-connect/src/lib.rs

@@ -30,6 +30,7 @@ use crate::session::Session;
 mod config;
 mod err;
 mod op;
+
 mod session;
 mod translation;
 pub mod util;

src/daft-connect/src/op/execute/root.rs

@@ -10,7 +10,6 @@ use crate::{
     op::execute::{ExecuteStream, PlanIds},
     session::Session,
     translation,
-    translation::Plan,
 };
 
 impl Session {
@@ -30,19 +29,25 @@ impl Session {
         let finished = context.finished();
 
         let (tx, rx) = tokio::sync::mpsc::channel::<eyre::Result<ExecutePlanResponse>>(1);
+
+        let pset = self.pset.clone();
+
         tokio::spawn(async move {
             let execution_fut = async {
-                let Plan { builder, psets } = translation::to_logical_plan(command).await?;
+                let translator = translation::SparkAnalyzer::new(&pset);
+                let lp = translator.to_logical_plan(command).await?;
 
                 // todo: convert optimize to async (looks like A LOT of work)... it touches a lot of API
                 // I tried and spent about an hour and gave up ~ Andrew Gazelka 🪦 2024-12-09
-                let optimized_plan = tokio::task::spawn_blocking(move || builder.optimize())
+                let optimized_plan = tokio::task::spawn_blocking(move || lp.optimize())
                     .await
                     .unwrap()?;
 
                 let cfg = Arc::new(DaftExecutionConfig::default());
                 let native_executor = NativeExecutor::from_logical_plan_builder(&optimized_plan)?;
-                let mut result_stream = native_executor.run(psets, cfg, None)?.into_stream();
+
+                let mut result_stream =
+                    native_executor.run(pset.as_ref(), cfg, None)?.into_stream();
 
                 while let Some(result) = result_stream.next().await {
                     let result = result?;

src/daft-connect/src/op/execute/write.rs

@@ -32,6 +32,7 @@ impl Session {
         };
 
         let finished = context.finished();
+        let pset = self.pset.clone();
 
         let result = async move {
             let WriteOperation {
@@ -109,18 +110,20 @@ impl Session {
                 }
             };
 
-            let mut plan = translation::to_logical_plan(input).await?;
+            let translator = translation::SparkAnalyzer::new(&pset);
 
-            plan.builder = plan
-                .builder
+            let plan = translator.to_logical_plan(input).await?;
+
+            let plan = plan
                 .table_write(&path, FileFormat::Parquet, None, None, None)
                 .wrap_err("Failed to create table write plan")?;
 
-            let optimized_plan = plan.builder.optimize()?;
+            let optimized_plan = plan.optimize()?;
             let cfg = DaftExecutionConfig::default();
             let native_executor = NativeExecutor::from_logical_plan_builder(&optimized_plan)?;
+
             let mut result_stream = native_executor
-                .run(plan.psets, cfg.into(), None)?
+                .run(pset.as_ref(), cfg.into(), None)?
                 .into_stream();
 
             // this is so we make sure the operation is actually done