rename python function to match rust

apache · Dec 14, 2024 · 408a43f · 408a43f
1 parent f73c5ce
commit 408a43f
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/datafusion_ray/context.py b/datafusion_ray/context.py
@@ -50,7 +50,7 @@ def execute_query_stage(
 
     # if the query stage has a single output partition then we need to execute for the output
     # partition, otherwise we need to execute in parallel for each input partition
-    concurrency = stage.get_input_partition_count()
+    concurrency = stage.get_execution_partition_count()
     output_partitions_count = stage.get_output_partition_count()
     if output_partitions_count == 1:
         # reduce stage

diff --git a/src/query_stage.rs b/src/query_stage.rs
@@ -60,7 +60,7 @@ impl PyQueryStage {
         self.stage.get_child_stage_ids()
     }
 
-    pub fn get_input_partition_count(&self) -> usize {
+    pub fn get_execution_partition_count(&self) -> usize {
         self.stage.get_execution_partition_count()
     }