refactor

Author: adriangb

Cargo.lock

@@ -86,6 +86,7 @@ log = { workspace = true }
 object_store = { workspace = true, optional = true }
 parquet = { workspace = true, optional = true, default-features = true }
 recursive = { workspace = true, optional = true }
+strum = { workspace = true, features = ["derive"] }
 sqlparser = { workspace = true, optional = true }
 tokio = { workspace = true }
 

datafusion/common/Cargo.toml

@@ -15,11 +15,14 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::collections::HashSet;
 use std::fmt::{self, Display};
 use std::str::FromStr;
 
 use arrow::compute::CastOptions;
 use arrow::util::display::{DurationFormat, FormatOptions};
+use itertools::Itertools;
+use strum::EnumString;
 
 use crate::config::{ConfigField, Visit};
 use crate::error::{DataFusionError, Result};
@@ -222,14 +225,26 @@ impl ConfigField for ExplainFormat {
 /// (when specified) or aggregated metrics are displayed (when omitted).
 /// In contrast, `MetricType` determines which *levels* of metrics are
 /// displayed.
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, strum::Display, EnumString)]
+#[strum(
+    serialize_all = "lowercase",
+    ascii_case_insensitive,
+    parse_err_ty = DataFusionError,
+    parse_err_fn = metric_type_parse_error
+)]
 pub enum MetricType {
     /// Common metrics for high-level insights (answering which operator is slow)
     Summary,
     /// For deep operator-level introspection for developers
     Dev,
 }
 
+fn metric_type_parse_error(s: &str) -> DataFusionError {
+    DataFusionError::Configuration(format!(
+        "Invalid explain analyze level '{s}'. Expected 'summary' or 'dev'."
+    ))
+}
+
 impl MetricType {
     /// Returns the set of metric types that should be shown for this level.
     ///
@@ -243,36 +258,13 @@ impl MetricType {
     }
 }
 
-impl FromStr for MetricType {
-    type Err = DataFusionError;
-
-    fn from_str(s: &str) -> Result<Self, Self::Err> {
-        match s.trim().to_lowercase().as_str() {
-            "summary" => Ok(Self::Summary),
-            "dev" => Ok(Self::Dev),
-            other => Err(DataFusionError::Configuration(format!(
-                "Invalid explain analyze level. Expected 'summary' or 'dev'. Got '{other}'"
-            ))),
-        }
-    }
-}
-
-impl Display for MetricType {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        match self {
-            Self::Summary => write!(f, "summary"),
-            Self::Dev => write!(f, "dev"),
-        }
-    }
-}
-
 impl ConfigField for MetricType {
     fn visit<V: Visit>(&self, v: &mut V, key: &str, description: &'static str) {
         v.some(key, self, description)
     }
 
     fn set(&mut self, _: &str, value: &str) -> Result<()> {
-        *self = MetricType::from_str(value)?;
+        *self = value.parse()?;
         Ok(())
     }
 }
@@ -301,7 +293,13 @@ impl ConfigField for MetricType {
 /// Metrics that do not declare a category (the default for custom
 /// `Count` / `Gauge` metrics) are treated as
 /// [`Uncategorized`](Self::Uncategorized) for filtering purposes.
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, strum::Display, EnumString)]
+#[strum(
+    serialize_all = "lowercase",
+    ascii_case_insensitive,
+    parse_err_ty = DataFusionError,
+    parse_err_fn = metric_category_parse_error
+)]
 pub enum MetricCategory {
     /// Row counts and related dimensionless counters: `output_rows`,
     /// `spilled_rows`, `output_batches`, pruning metrics, ratios, etc.
@@ -333,32 +331,11 @@ pub enum MetricCategory {
     Uncategorized,
 }
 
-impl FromStr for MetricCategory {
-    type Err = DataFusionError;
-
-    fn from_str(s: &str) -> Result<Self, Self::Err> {
-        match s.trim().to_lowercase().as_str() {
-            "rows" => Ok(Self::Rows),
-            "bytes" => Ok(Self::Bytes),
-            "timing" => Ok(Self::Timing),
-            "uncategorized" => Ok(Self::Uncategorized),
-            other => Err(DataFusionError::Configuration(format!(
-                "Invalid metric category '{other}'. \
-                 Expected 'rows', 'bytes', 'timing', or 'uncategorized'."
-            ))),
-        }
-    }
-}
-
-impl Display for MetricCategory {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        match self {
-            Self::Rows => write!(f, "rows"),
-            Self::Bytes => write!(f, "bytes"),
-            Self::Timing => write!(f, "timing"),
-            Self::Uncategorized => write!(f, "uncategorized"),
-        }
-    }
+fn metric_category_parse_error(s: &str) -> DataFusionError {
+    DataFusionError::Configuration(format!(
+        "Invalid metric category '{s}'. \
+         Expected 'rows', 'bytes', 'timing', or 'uncategorized'."
+    ))
 }
 
 /// Controls which [`MetricCategory`] values are shown in `EXPLAIN ANALYZE`.
@@ -389,12 +366,11 @@ impl FromStr for ExplainAnalyzeCategories {
             "all" => Ok(Self::All),
             "none" => Ok(Self::Only(vec![])),
             other => {
-                let mut cats = Vec::new();
+                let mut cats = HashSet::new();
                 for part in other.split(',') {
-                    cats.push(part.trim().parse::<MetricCategory>()?);
+                    cats.insert(part.trim().parse::<MetricCategory>()?);
                 }
-                cats.dedup();
-                Ok(Self::Only(cats))
+                Ok(Self::Only(cats.into_iter().collect_vec()))
             }
         }
     }