[Minor] Use std::thread::available_parallelism instead of num_cpus (#…

…13579)

* Use std::thread::available_parallelism

* Use std::thread::available_parallelism

* Use std::thread::available_parallelism

* Use std::thread::available_parallelism

* Use std::thread::available_parallelism

* Use std::thread::available_parallelism

Cargo.toml

@@ -129,7 +129,6 @@ hashbrown = { version = "0.14.5", features = ["raw"] }
 indexmap = "2.0.0"
 itertools = "0.13"
 log = "^0.4"
-num_cpus = "1.13.0"
 object_store = { version = "0.11.0", default-features = false }
 parking_lot = "0.12"
 parquet = { version = "53.3.0", default-features = false, features = [

benchmarks/Cargo.toml

@@ -42,7 +42,6 @@ env_logger = { workspace = true }
 futures = { workspace = true }
 log = { workspace = true }
 mimalloc = { version = "0.1", optional = true, default-features = false }
-num_cpus = { workspace = true }
 parquet = { workspace = true, default-features = true }
 serde = { version = "1.0.136", features = ["derive"] }
 serde_json = { workspace = true }

benchmarks/src/bin/external_aggr.rs

@@ -18,9 +18,11 @@
 //! external_aggr binary entrypoint
 
 use std::collections::HashMap;
+use std::num::NonZero;
 use std::path::PathBuf;
 use std::sync::Arc;
 use std::sync::OnceLock;
+use std::thread::available_parallelism;
 use structopt::StructOpt;
 
 use arrow::record_batch::RecordBatch;
@@ -325,7 +327,11 @@ impl ExternalAggrConfig {
     }
 
     fn partitions(&self) -> usize {
-        self.common.partitions.unwrap_or(num_cpus::get())
+        self.common.partitions.unwrap_or(
+            available_parallelism()
+                .unwrap_or(NonZero::new(1).unwrap())
+                .get(),
+        )
     }
 
     /// Parse memory limit from string to number of bytes

benchmarks/src/bin/h2o.rs

@@ -27,8 +27,10 @@ use datafusion::datasource::MemTable;
 use datafusion::prelude::CsvReadOptions;
 use datafusion::{arrow::util::pretty, error::Result, prelude::SessionContext};
 use datafusion_benchmarks::util::BenchmarkRun;
+use std::num::NonZero;
 use std::path::PathBuf;
 use std::sync::Arc;
+use std::thread::available_parallelism;
 use structopt::StructOpt;
 use tokio::time::Instant;
 
@@ -91,7 +93,9 @@ async fn group_by(opt: &GroupBy) -> Result<()> {
             .with_listing_options(ListingOptions::new(Arc::new(CsvFormat::default())))
             .with_schema(Arc::new(schema));
         let csv = ListingTable::try_new(listing_config)?;
-        let partition_size = num_cpus::get();
+        let partition_size = available_parallelism()
+            .unwrap_or(NonZero::new(1).unwrap())
+            .get();
         let memtable =
             MemTable::load(Arc::new(csv), Some(partition_size), &ctx.state()).await?;
         ctx.register_table("x", Arc::new(memtable))?;

benchmarks/src/imdb/run.rs

@@ -15,8 +15,10 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::num::NonZero;
 use std::path::PathBuf;
 use std::sync::Arc;
+use std::thread::available_parallelism;
 
 use super::{get_imdb_table_schema, get_query_sql, IMDB_TABLES};
 use crate::util::{BenchmarkRun, CommonOpt};
@@ -468,7 +470,11 @@ impl RunOpt {
     }
 
     fn partitions(&self) -> usize {
-        self.common.partitions.unwrap_or(num_cpus::get())
+        self.common.partitions.unwrap_or(
+            available_parallelism()
+                .unwrap_or(NonZero::new(1).unwrap())
+                .get(),
+        )
     }
 }
 

benchmarks/src/sort.rs

@@ -15,8 +15,10 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::num::NonZero;
 use std::path::PathBuf;
 use std::sync::Arc;
+use std::thread::available_parallelism;
 
 use crate::util::{AccessLogOpt, BenchmarkRun, CommonOpt};
 
@@ -147,7 +149,11 @@ impl RunOpt {
             rundata.start_new_case(title);
             for i in 0..self.common.iterations {
                 let config = SessionConfig::new().with_target_partitions(
-                    self.common.partitions.unwrap_or(num_cpus::get()),
+                    self.common.partitions.unwrap_or(
+                        available_parallelism()
+                            .unwrap_or(NonZero::new(1).unwrap())
+                            .get(),
+                    ),
                 );
                 let ctx = SessionContext::new_with_config(config);
                 let (rows, elapsed) =

benchmarks/src/sort_tpch.rs

@@ -22,8 +22,10 @@
 //! runs end-to-end sort queries and test the performance on multiple CPU cores.
 
 use futures::StreamExt;
+use std::num::NonZero;
 use std::path::PathBuf;
 use std::sync::Arc;
+use std::thread::available_parallelism;
 use structopt::StructOpt;
 
 use datafusion::datasource::file_format::parquet::ParquetFormat;
@@ -315,6 +317,10 @@ impl RunOpt {
     }
 
     fn partitions(&self) -> usize {
-        self.common.partitions.unwrap_or(num_cpus::get())
+        self.common.partitions.unwrap_or(
+            available_parallelism()
+                .unwrap_or(NonZero::new(1).unwrap())
+                .get(),
+        )
     }
 }

benchmarks/src/tpch/run.rs

@@ -15,8 +15,10 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::num::NonZero;
 use std::path::PathBuf;
 use std::sync::Arc;
+use std::thread::available_parallelism;
 
 use super::{
     get_query_sql, get_tbl_tpch_table_schema, get_tpch_table_schema, TPCH_TABLES,
@@ -296,7 +298,11 @@ impl RunOpt {
     }
 
     fn partitions(&self) -> usize {
-        self.common.partitions.unwrap_or(num_cpus::get())
+        self.common.partitions.unwrap_or(
+            available_parallelism()
+                .unwrap_or(NonZero::new(1).unwrap())
+                .get(),
+        )
     }
 }
 

benchmarks/src/util/options.rs

@@ -15,6 +15,8 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::{num::NonZero, thread::available_parallelism};
+
 use datafusion::prelude::SessionConfig;
 use structopt::StructOpt;
 
@@ -48,7 +50,13 @@ impl CommonOpt {
     /// Modify the existing config appropriately
     pub fn update_config(&self, config: SessionConfig) -> SessionConfig {
         config
-            .with_target_partitions(self.partitions.unwrap_or(num_cpus::get()))
+            .with_target_partitions(
+                self.partitions.unwrap_or(
+                    available_parallelism()
+                        .unwrap_or(NonZero::new(1).unwrap())
+                        .get(),
+                ),
+            )
             .with_batch_size(self.batch_size)
     }
 }

benchmarks/src/util/run.rs

@@ -20,7 +20,9 @@ use serde::{Serialize, Serializer};
 use serde_json::Value;
 use std::{
     collections::HashMap,
+    num::NonZero,
     path::Path,
+    thread::available_parallelism,
     time::{Duration, SystemTime},
 };
 
@@ -68,7 +70,9 @@ impl RunContext {
         Self {
             benchmark_version: env!("CARGO_PKG_VERSION").to_owned(),
             datafusion_version: DATAFUSION_VERSION.to_owned(),
-            num_cpus: num_cpus::get(),
+            num_cpus: available_parallelism()
+                .unwrap_or(NonZero::new(1).unwrap())
+                .get(),
             start_time: SystemTime::now(),
             arguments: std::env::args().skip(1).collect::<Vec<String>>(),
         }

datafusion/common/Cargo.toml

@@ -58,7 +58,6 @@ half = { workspace = true }
 hashbrown = { workspace = true }
 indexmap = { workspace = true }
 libc = "0.2.140"
-num_cpus = { workspace = true }
 object_store = { workspace = true, optional = true }
 parquet = { workspace = true, optional = true, default-features = true }
 paste = "1.0.15"

datafusion/common/src/config.rs

@@ -20,7 +20,9 @@
 use std::any::Any;
 use std::collections::{BTreeMap, HashMap};
 use std::fmt::{self, Display};
+use std::num::NonZero;
 use std::str::FromStr;
+use std::thread::available_parallelism;
 
 use crate::error::_config_err;
 use crate::parsers::CompressionTypeVariant;
@@ -250,7 +252,7 @@ config_namespace! {
         /// concurrency.
         ///
         /// Defaults to the number of CPU cores on the system
-        pub target_partitions: usize, default = num_cpus::get()
+        pub target_partitions: usize, default = available_parallelism().unwrap_or(NonZero::new(1).unwrap()).get()
 
         /// The default time zone
         ///
@@ -266,7 +268,7 @@ config_namespace! {
         /// This is mostly use to plan `UNION` children in parallel.
         ///
         /// Defaults to the number of CPU cores on the system
-        pub planning_concurrency: usize, default = num_cpus::get()
+        pub planning_concurrency: usize, default = available_parallelism().unwrap_or(NonZero::new(1).unwrap()).get()
 
         /// When set to true, skips verifying that the schema produced by
         /// planning the input of `LogicalPlan::Aggregate` exactly matches the

datafusion/core/Cargo.toml

@@ -116,7 +116,6 @@ glob = "0.3.0"
 itertools = { workspace = true }
 log = { workspace = true }
 num-traits = { version = "0.2", optional = true }
-num_cpus = { workspace = true }
 object_store = { workspace = true }
 parking_lot = { workspace = true }
 parquet = { workspace = true, optional = true, default-features = true }

datafusion/core/tests/fuzz_cases/aggregation_fuzzer/context_generator.rs

@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use std::{cmp, sync::Arc};
+use std::{cmp, num::NonZero, sync::Arc, thread::available_parallelism};
 
 use datafusion::{
     datasource::MemTable,
@@ -73,7 +73,9 @@ impl SessionContextGenerator {
         ];
 
         let max_batch_size = cmp::max(1, dataset_ref.total_rows_num);
-        let max_target_partitions = num_cpus::get();
+        let max_target_partitions = available_parallelism()
+            .unwrap_or(NonZero::new(1).unwrap())
+            .get();
 
         Self {
             dataset: dataset_ref,

datafusion/core/tests/sql/mod.rs

@@ -15,7 +15,9 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::num::NonZero;
 use std::sync::Arc;
+use std::thread::available_parallelism;
 
 use arrow::{
     array::*, datatypes::*, record_batch::RecordBatch,
@@ -259,7 +261,12 @@ impl ExplainNormalizer {
 
         // convert things like partitioning=RoundRobinBatch(16)
         // to partitioning=RoundRobinBatch(NUM_CORES)
-        let needle = format!("RoundRobinBatch({})", num_cpus::get());
+        let needle = format!(
+            "RoundRobinBatch({})",
+            available_parallelism()
+                .unwrap_or(NonZero::new(1).unwrap())
+                .get()
+        );
         replacements.push((needle, "RoundRobinBatch(NUM_CORES)".to_string()));
 
         Self { replacements }

datafusion/sqllogictest/Cargo.toml

@@ -70,7 +70,6 @@ postgres = [
 
 [dev-dependencies]
 env_logger = { workspace = true }
-num_cpus = { workspace = true }
 tokio = { workspace = true, features = ["rt-multi-thread"] }
 
 [[test]]

datafusion/sqllogictest/bin/sqllogictests.rs

@@ -17,7 +17,9 @@
 
 use std::ffi::OsStr;
 use std::fs;
+use std::num::NonZero;
 use std::path::{Path, PathBuf};
+use std::thread::available_parallelism;
 
 use clap::Parser;
 use datafusion_sqllogictest::{DataFusion, TestContext};
@@ -112,7 +114,11 @@ async fn run_tests() -> Result<()> {
             .join()
         })
         // run up to num_cpus streams in parallel
-        .buffer_unordered(num_cpus::get())
+        .buffer_unordered(
+            available_parallelism()
+                .unwrap_or(NonZero::new(1).unwrap())
+                .get(),
+        )
         .flat_map(|result| {
             // Filter out any Ok() leaving only the DataFusionErrors
             futures::stream::iter(match result {