| 219 | } |
| 220 | |
| 221 | fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> { |
| 222 | assert!( |
| 223 | args.args.len() == 2, |
| 224 | "json_get_str requires exactly 2 arguments" |
| 225 | ); |
| 226 | let key = match &args.args[0] { |
| 227 | ColumnarValue::Scalar(ScalarValue::Utf8(Some(key))) => key, |
| 228 | _ => { |
| 229 | return Err(exec_datafusion_err!( |
| 230 | "json_get_str first argument must be a string" |
| 231 | )); |
| 232 | } |
| 233 | }; |
| 234 | // We expect a string array that contains JSON strings |
| 235 | let json_array = match &args.args[1] { |
| 236 | ColumnarValue::Array(array) => array |
| 237 | .as_any() |
| 238 | .downcast_ref::<StringArray>() |
| 239 | .ok_or_else(|| { |
| 240 | exec_datafusion_err!( |
| 241 | "json_get_str second argument must be a string array" |
| 242 | ) |
| 243 | })?, |
| 244 | _ => { |
| 245 | return Err(exec_datafusion_err!( |
| 246 | "json_get_str second argument must be a string array" |
| 247 | )); |
| 248 | } |
| 249 | }; |
| 250 | let values = json_array |
| 251 | .iter() |
| 252 | .map(|value| { |
| 253 | value.and_then(|v| { |
| 254 | let json_value: serde_json::Value = |
| 255 | serde_json::from_str(v).unwrap_or_default(); |
| 256 | json_value.get(key).map(|v| v.to_string()) |
| 257 | }) |
| 258 | }) |
| 259 | .collect::<StringArray>(); |
| 260 | Ok(ColumnarValue::Array(Arc::new(values))) |
| 261 | } |
| 262 | } |
| 263 | |
| 264 | /// Factory for creating ShreddedJsonRewriter instances |