fixup! [HSTACK] Building blocks for Ray DataFusionDatasource

Author: ccciudatu

python/datafusion/__init__.py

@@ -29,7 +29,7 @@
 from . import functions, object_store, substrait
 
 # The following imports are okay to remain as opaque to the user.
-from ._internal import Config
+from ._internal import Config, partition_stream
 from .catalog import Catalog, Database, Table
 from .common import (
     DFSchema,
@@ -81,6 +81,7 @@
     "functions",
     "object_store",
     "substrait",
+    "partition_stream",
 ]
 
 

python/datafusion/dataframe.py

@@ -797,8 +797,8 @@ def count(self) -> int:
         """
         return self.df.count()
 
-    def distributed_plan(self, num_shards: int):
-        return self.df.distributed_plan(num_shards)
+    def distributed_plan(self):
+        return self.df.distributed_plan()
 
     @deprecated("Use :py:func:`unnest_columns` instead.")
     def unnest_column(self, column: str, preserve_nulls: bool = True) -> DataFrame:

src/dataframe.rs

@@ -29,7 +29,7 @@ use datafusion::arrow::pyarrow::{PyArrowType, ToPyArrow};
 use datafusion::arrow::util::pretty;
 use datafusion::common::stats::Precision;
 use datafusion::common::{DFSchema, DataFusionError, UnnestOptions};
-use datafusion::config::{ConfigOptions, CsvOptions, TableParquetOptions};
+use datafusion::config::{CsvOptions, TableParquetOptions};
 use datafusion::dataframe::{DataFrame, DataFrameWriteOptions};
 use datafusion::execution::runtime_env::RuntimeEnvBuilder;
 use datafusion::execution::SendableRecordBatchStream;
@@ -669,10 +669,10 @@ impl PyDataFrame {
         Ok(wait_for_future(py, self.df.as_ref().clone().count())?)
     }
 
-    fn distributed_plan(&self, parallelism: usize, py: Python<'_>) -> PyResult<DistributedPlan> {
+    fn distributed_plan(&self, py: Python<'_>) -> PyResult<DistributedPlan> {
         let future_plan = self.df.as_ref().clone().create_physical_plan();
         let physical_plan = wait_for_future(py, future_plan).map_err(py_datafusion_err)?;
-        DistributedPlan::try_new(physical_plan, parallelism).map_err(py_datafusion_err)
+        DistributedPlan::try_new(physical_plan).map_err(py_datafusion_err)
     }
 
 }
@@ -682,7 +682,7 @@ impl PyDataFrame {
 pub struct DistributedPlan {
     physical_plan: Vec<u8>,
     schema: PyDFSchema,
-    partitions: usize,
+    num_partitions: usize,
     num_bytes: Option<usize>,
     num_rows: Option<usize>,
 }
@@ -693,29 +693,27 @@ fn codec() -> &'static dyn PhysicalExtensionCodec {
 }
 
 impl DistributedPlan {
-    fn try_new(plan: Arc<dyn ExecutionPlan>, parallelism: usize) -> Result<Self, DataFusionError> {
+    fn try_new(plan: Arc<dyn ExecutionPlan>) -> Result<Self, DataFusionError> {
         fn extract(prec: Precision<usize>) -> Option<usize> {
             match prec {
                 Precision::Exact(n) => Some(n),
                 _ => None,
             }
         }
         let (num_bytes, num_rows) = if let Ok(stats) = plan.statistics() {
-            let num_bytes = extract(stats.total_byte_size);
-            let num_rows = extract(stats.num_rows);
-            (num_bytes, num_rows)
+            let bytes = extract(stats.total_byte_size);
+            let rows = extract(stats.num_rows);
+            (bytes, rows)
         } else {
             (None, None)
         };
 
         let schema = DFSchema::try_from(plan.schema())
             .map(PyDFSchema::from)?;
-        let plan = plan.repartitioned(parallelism, &ConfigOptions::default())?
-            .unwrap_or(plan);
-        let partitions = plan.properties().partitioning.partition_count();
+        let num_partitions = plan.properties().partitioning.partition_count();
         let physical_plan = PhysicalPlanNode::try_from_physical_plan(plan, codec())?
             .encode_to_vec();
-        Ok(Self { physical_plan, schema, partitions, num_bytes, num_rows })
+        Ok(Self { physical_plan, schema, num_partitions, num_bytes, num_rows })
     }
 
 }