Skip to content

Commit dc6a712

Browse files
cj-zhukovSergey Zhukovalamb
authored
Use strum in the examples (#19126) (#19205)
## Which issue does this PR close? <!-- We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123. --> - Closes ##19126. ## Rationale for this change <!-- Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed. Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes. --> ## What changes are included in this PR? <!-- There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR. --> ## Are these changes tested? <!-- We typically require tests for all PRs in order to: 1. Prevent the code from being accidentally broken by subsequent changes 2. Serve as another way to document the expected behavior of the code If tests are not included in your PR, please explain why (for example, are they covered by existing tests)? --> ## Are there any user-facing changes? <!-- If there are user-facing changes then we may require documentation to be updated before approving the PR. --> <!-- If there are any breaking changes to public APIs, please add the `api change` label. --> --------- Co-authored-by: Sergey Zhukov <[email protected]> Co-authored-by: Andrew Lamb <[email protected]>
1 parent cd160de commit dc6a712

File tree

15 files changed

+199
-679
lines changed

15 files changed

+199
-679
lines changed

‎Cargo.lock‎

Lines changed: 2 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎Cargo.toml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -182,6 +182,8 @@ regex = "1.12"
182182
rstest = "0.26.1"
183183
serde_json = "1"
184184
sqlparser = { version = "0.59.0", default-features = false, features = ["std", "visitor"] }
185+
strum = "0.27.2"
186+
strum_macros = "0.27.2"
185187
tempfile = "3"
186188
testcontainers = { version = "0.25.2", features = ["default"] }
187189
testcontainers-modules = { version = "0.13" }

‎datafusion-examples/Cargo.toml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,8 @@ object_store = { workspace = true, features = ["aws", "http"] }
6060
prost = { workspace = true }
6161
rand = { workspace = true }
6262
serde_json = { workspace = true }
63+
strum = { workspace = true }
64+
strum_macros = { workspace = true }
6365
tempfile = { workspace = true }
6466
test-utils = { path = "../test-utils" }
6567
tokio = { workspace = true, features = ["rt-multi-thread", "parking_lot"] }

‎datafusion-examples/examples/builtin_functions/main.rs‎

Lines changed: 14 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -34,67 +34,31 @@ mod date_time;
3434
mod function_factory;
3535
mod regexp;
3636

37-
use std::str::FromStr;
38-
3937
use datafusion::error::{DataFusionError, Result};
38+
use strum::{IntoEnumIterator, VariantNames};
39+
use strum_macros::{Display, EnumIter, EnumString, VariantNames};
4040

41+
#[derive(EnumIter, EnumString, Display, VariantNames)]
42+
#[strum(serialize_all = "snake_case")]
4143
enum ExampleKind {
4244
All,
4345
DateTime,
4446
FunctionFactory,
4547
Regexp,
4648
}
4749

48-
impl AsRef<str> for ExampleKind {
49-
fn as_ref(&self) -> &str {
50-
match self {
51-
Self::All => "all",
52-
Self::DateTime => "date_time",
53-
Self::FunctionFactory => "function_factory",
54-
Self::Regexp => "regexp",
55-
}
56-
}
57-
}
58-
59-
impl FromStr for ExampleKind {
60-
type Err = DataFusionError;
61-
62-
fn from_str(s: &str) -> Result<Self> {
63-
match s {
64-
"all" => Ok(Self::All),
65-
"date_time" => Ok(Self::DateTime),
66-
"function_factory" => Ok(Self::FunctionFactory),
67-
"regexp" => Ok(Self::Regexp),
68-
_ => Err(DataFusionError::Execution(format!("Unknown example: {s}"))),
69-
}
70-
}
71-
}
72-
7350
impl ExampleKind {
74-
const ALL_VARIANTS: [Self; 4] = [
75-
Self::All,
76-
Self::DateTime,
77-
Self::FunctionFactory,
78-
Self::Regexp,
79-
];
80-
81-
const RUNNABLE_VARIANTS: [Self; 3] =
82-
[Self::DateTime, Self::FunctionFactory, Self::Regexp];
83-
8451
const EXAMPLE_NAME: &str = "builtin_functions";
8552

86-
fn variants() -> Vec<&'static str> {
87-
Self::ALL_VARIANTS
88-
.iter()
89-
.map(|example| example.as_ref())
90-
.collect()
53+
fn runnable() -> impl Iterator<Item = ExampleKind> {
54+
ExampleKind::iter().filter(|v| !matches!(v, ExampleKind::All))
9155
}
9256

9357
async fn run(&self) -> Result<()> {
9458
match self {
9559
ExampleKind::All => {
96-
for example in ExampleKind::RUNNABLE_VARIANTS {
97-
println!("Running example: {}", example.as_ref());
60+
for example in ExampleKind::runnable() {
61+
println!("Running example: {example}");
9862
Box::pin(example.run()).await?;
9963
}
10064
}
@@ -111,14 +75,14 @@ async fn main() -> Result<()> {
11175
let usage = format!(
11276
"Usage: cargo run --example {} -- [{}]",
11377
ExampleKind::EXAMPLE_NAME,
114-
ExampleKind::variants().join("|")
78+
ExampleKind::VARIANTS.join("|")
11579
);
11680

117-
let arg = std::env::args().nth(1).ok_or_else(|| {
118-
eprintln!("{usage}");
119-
DataFusionError::Execution("Missing argument".to_string())
120-
})?;
81+
let example: ExampleKind = std::env::args()
82+
.nth(1)
83+
.ok_or_else(|| DataFusionError::Execution(format!("Missing argument. {usage}")))?
84+
.parse()
85+
.map_err(|_| DataFusionError::Execution(format!("Unknown example. {usage}")))?;
12186

122-
let example = arg.parse::<ExampleKind>()?;
12387
example.run().await
12488
}

‎datafusion-examples/examples/custom_data_source/main.rs‎

Lines changed: 14 additions & 69 deletions
Original file line numberDiff line numberDiff line change
@@ -42,10 +42,12 @@ mod custom_file_format;
4242
mod default_column_values;
4343
mod file_stream_provider;
4444

45-
use std::str::FromStr;
46-
4745
use datafusion::error::{DataFusionError, Result};
46+
use strum::{IntoEnumIterator, VariantNames};
47+
use strum_macros::{Display, EnumIter, EnumString, VariantNames};
4848

49+
#[derive(EnumIter, EnumString, Display, VariantNames)]
50+
#[strum(serialize_all = "snake_case")]
4951
enum ExampleKind {
5052
All,
5153
CsvJsonOpener,
@@ -57,75 +59,18 @@ enum ExampleKind {
5759
FileStreamProvider,
5860
}
5961

60-
impl AsRef<str> for ExampleKind {
61-
fn as_ref(&self) -> &str {
62-
match self {
63-
Self::All => "all",
64-
Self::CsvJsonOpener => "csv_json_opener",
65-
Self::CsvSqlStreaming => "csv_sql_streaming",
66-
Self::CustomDatasource => "custom_datasource",
67-
Self::CustomFileCasts => "custom_file_casts",
68-
Self::CustomFileFormat => "custom_file_format",
69-
Self::DefaultColumnValues => "default_column_values",
70-
Self::FileStreamProvider => "file_stream_provider",
71-
}
72-
}
73-
}
74-
75-
impl FromStr for ExampleKind {
76-
type Err = DataFusionError;
77-
78-
fn from_str(s: &str) -> Result<Self> {
79-
match s {
80-
"all" => Ok(Self::All),
81-
"csv_json_opener" => Ok(Self::CsvJsonOpener),
82-
"csv_sql_streaming" => Ok(Self::CsvSqlStreaming),
83-
"custom_datasource" => Ok(Self::CustomDatasource),
84-
"custom_file_casts" => Ok(Self::CustomFileCasts),
85-
"custom_file_format" => Ok(Self::CustomFileFormat),
86-
"default_column_values" => Ok(Self::DefaultColumnValues),
87-
"file_stream_provider" => Ok(Self::FileStreamProvider),
88-
_ => Err(DataFusionError::Execution(format!("Unknown example: {s}"))),
89-
}
90-
}
91-
}
92-
9362
impl ExampleKind {
94-
const ALL_VARIANTS: [Self; 8] = [
95-
Self::All,
96-
Self::CsvJsonOpener,
97-
Self::CsvSqlStreaming,
98-
Self::CustomDatasource,
99-
Self::CustomFileCasts,
100-
Self::CustomFileFormat,
101-
Self::DefaultColumnValues,
102-
Self::FileStreamProvider,
103-
];
104-
105-
const RUNNABLE_VARIANTS: [Self; 7] = [
106-
Self::CsvJsonOpener,
107-
Self::CsvSqlStreaming,
108-
Self::CustomDatasource,
109-
Self::CustomFileCasts,
110-
Self::CustomFileFormat,
111-
Self::DefaultColumnValues,
112-
Self::FileStreamProvider,
113-
];
114-
11563
const EXAMPLE_NAME: &str = "custom_data_source";
11664

117-
fn variants() -> Vec<&'static str> {
118-
Self::ALL_VARIANTS
119-
.iter()
120-
.map(|example| example.as_ref())
121-
.collect()
65+
fn runnable() -> impl Iterator<Item = ExampleKind> {
66+
ExampleKind::iter().filter(|v| !matches!(v, ExampleKind::All))
12267
}
12368

12469
async fn run(&self) -> Result<()> {
12570
match self {
12671
ExampleKind::All => {
127-
for example in ExampleKind::RUNNABLE_VARIANTS {
128-
println!("Running example: {}", example.as_ref());
72+
for example in ExampleKind::runnable() {
73+
println!("Running example: {example}");
12974
Box::pin(example.run()).await?;
13075
}
13176
}
@@ -158,14 +103,14 @@ async fn main() -> Result<()> {
158103
let usage = format!(
159104
"Usage: cargo run --example {} -- [{}]",
160105
ExampleKind::EXAMPLE_NAME,
161-
ExampleKind::variants().join("|")
106+
ExampleKind::VARIANTS.join("|")
162107
);
163108

164-
let arg = std::env::args().nth(1).ok_or_else(|| {
165-
eprintln!("{usage}");
166-
DataFusionError::Execution("Missing argument".to_string())
167-
})?;
109+
let example: ExampleKind = std::env::args()
110+
.nth(1)
111+
.ok_or_else(|| DataFusionError::Execution(format!("Missing argument. {usage}")))?
112+
.parse()
113+
.map_err(|_| DataFusionError::Execution(format!("Unknown example. {usage}")))?;
168114

169-
let example = arg.parse::<ExampleKind>()?;
170115
example.run().await
171116
}

‎datafusion-examples/examples/data_io/main.rs‎

Lines changed: 14 additions & 81 deletions
Original file line numberDiff line numberDiff line change
@@ -48,10 +48,12 @@ mod parquet_index;
4848
mod query_http_csv;
4949
mod remote_catalog;
5050

51-
use std::str::FromStr;
52-
5351
use datafusion::error::{DataFusionError, Result};
52+
use strum::{IntoEnumIterator, VariantNames};
53+
use strum_macros::{Display, EnumIter, EnumString, VariantNames};
5454

55+
#[derive(EnumIter, EnumString, Display, VariantNames)]
56+
#[strum(serialize_all = "snake_case")]
5557
enum ExampleKind {
5658
All,
5759
Catalog,
@@ -66,87 +68,18 @@ enum ExampleKind {
6668
RemoteCatalog,
6769
}
6870

69-
impl AsRef<str> for ExampleKind {
70-
fn as_ref(&self) -> &str {
71-
match self {
72-
Self::All => "all",
73-
Self::Catalog => "catalog",
74-
Self::JsonShredding => "json_shredding",
75-
Self::ParquetAdvIdx => "parquet_adv_idx",
76-
Self::ParquetEmbIdx => "parquet_emb_idx",
77-
Self::ParquetEnc => "parquet_enc",
78-
Self::ParquetEncWithKms => "parquet_enc_with_kms",
79-
Self::ParquetExecVisitor => "parquet_exec_visitor",
80-
Self::ParquetIdx => "parquet_idx",
81-
Self::QueryHttpCsv => "query_http_csv",
82-
Self::RemoteCatalog => "remote_catalog",
83-
}
84-
}
85-
}
86-
87-
impl FromStr for ExampleKind {
88-
type Err = DataFusionError;
89-
90-
fn from_str(s: &str) -> Result<Self> {
91-
match s {
92-
"all" => Ok(Self::All),
93-
"catalog" => Ok(Self::Catalog),
94-
"json_shredding" => Ok(Self::JsonShredding),
95-
"parquet_adv_idx" => Ok(Self::ParquetAdvIdx),
96-
"parquet_emb_idx" => Ok(Self::ParquetEmbIdx),
97-
"parquet_enc" => Ok(Self::ParquetEnc),
98-
"parquet_enc_with_kms" => Ok(Self::ParquetEncWithKms),
99-
"parquet_exec_visitor" => Ok(Self::ParquetExecVisitor),
100-
"parquet_idx" => Ok(Self::ParquetIdx),
101-
"query_http_csv" => Ok(Self::QueryHttpCsv),
102-
"remote_catalog" => Ok(Self::RemoteCatalog),
103-
_ => Err(DataFusionError::Execution(format!("Unknown example: {s}"))),
104-
}
105-
}
106-
}
107-
10871
impl ExampleKind {
109-
const ALL_VARIANTS: [Self; 11] = [
110-
Self::All,
111-
Self::Catalog,
112-
Self::JsonShredding,
113-
Self::ParquetAdvIdx,
114-
Self::ParquetEmbIdx,
115-
Self::ParquetEnc,
116-
Self::ParquetEncWithKms,
117-
Self::ParquetExecVisitor,
118-
Self::ParquetIdx,
119-
Self::QueryHttpCsv,
120-
Self::RemoteCatalog,
121-
];
122-
123-
const RUNNABLE_VARIANTS: [Self; 10] = [
124-
Self::Catalog,
125-
Self::JsonShredding,
126-
Self::ParquetAdvIdx,
127-
Self::ParquetEmbIdx,
128-
Self::ParquetEnc,
129-
Self::ParquetEncWithKms,
130-
Self::ParquetExecVisitor,
131-
Self::ParquetIdx,
132-
Self::QueryHttpCsv,
133-
Self::RemoteCatalog,
134-
];
135-
13672
const EXAMPLE_NAME: &str = "data_io";
13773

138-
fn variants() -> Vec<&'static str> {
139-
Self::ALL_VARIANTS
140-
.iter()
141-
.map(|example| example.as_ref())
142-
.collect()
74+
fn runnable() -> impl Iterator<Item = ExampleKind> {
75+
ExampleKind::iter().filter(|v| !matches!(v, ExampleKind::All))
14376
}
14477

14578
async fn run(&self) -> Result<()> {
14679
match self {
14780
ExampleKind::All => {
148-
for example in ExampleKind::RUNNABLE_VARIANTS {
149-
println!("Running example: {}", example.as_ref());
81+
for example in ExampleKind::runnable() {
82+
println!("Running example: {example}");
15083
Box::pin(example.run()).await?;
15184
}
15285
}
@@ -178,14 +111,14 @@ async fn main() -> Result<()> {
178111
let usage = format!(
179112
"Usage: cargo run --example {} -- [{}]",
180113
ExampleKind::EXAMPLE_NAME,
181-
ExampleKind::variants().join("|")
114+
ExampleKind::VARIANTS.join("|")
182115
);
183116

184-
let arg = std::env::args().nth(1).ok_or_else(|| {
185-
eprintln!("{usage}");
186-
DataFusionError::Execution("Missing argument".to_string())
187-
})?;
117+
let example: ExampleKind = std::env::args()
118+
.nth(1)
119+
.ok_or_else(|| DataFusionError::Execution(format!("Missing argument. {usage}")))?
120+
.parse()
121+
.map_err(|_| DataFusionError::Execution(format!("Unknown example. {usage}")))?;
188122

189-
let example = arg.parse::<ExampleKind>()?;
190123
example.run().await
191124
}

0 commit comments

Comments
 (0)