use arrow::datatypes::{DataType, Field};
use datafusion_common::{DataFusionError, Result, ScalarValue};
pub fn get_indexed_field(data_type: &DataType, key: &ScalarValue) -> Result<Field> {
match (data_type, key) {
(DataType::List(lt), ScalarValue::Int64(Some(i))) => {
Ok(Field::new(i.to_string(), lt.data_type().clone(), true))
}
(DataType::Struct(fields), ScalarValue::Utf8(Some(s))) => {
if s.is_empty() {
Err(DataFusionError::Plan(
"Struct based indexed access requires a non empty string".to_string(),
))
} else {
let field = fields.iter().find(|f| f.name() == s);
match field {
None => Err(DataFusionError::Plan(format!(
"Field {s} not found in struct"
))),
Some(f) => Ok(f.clone()),
}
}
}
(DataType::Struct(_), _) => Err(DataFusionError::Plan(
"Only utf8 strings are valid as an indexed field in a struct".to_string(),
)),
(DataType::List(_), _) => Err(DataFusionError::Plan(
"Only ints are valid as an indexed field in a list".to_string(),
)),
_ => Err(DataFusionError::Plan(
"The expression to get an indexed field is only valid for `List` types"
.to_string(),
)),
}
}