dply/engine/
unnest.rs

1// Copyright (C) 2023 Vince Vasta
2// SPDX-License-Identifier: Apache-2.0
3//
4// Licensed under the Apache License, Version 2.0 (the "License");
5// you may not use this file except in compliance with the License.
6// You may obtain a copy of the License at
7//
8//     http://www.apache.org/licenses/LICENSE-2.0
9//
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15use anyhow::{bail, Result};
16use polars::prelude::*;
17
18use crate::parser::Expr;
19
20use super::*;
21
22/// Evaluates an unnest call.
23///
24/// Parameters are checked before evaluation by the typing module.
25pub fn eval(args: &[Expr], ctx: &mut Context) -> Result<()> {
26    if let Some(mut df) = ctx.take_df() {
27        for arg in args {
28            let column = args::identifier(arg);
29            let schema = df
30                .collect_schema()
31                .map_err(|e| anyhow!("unnest error: {e}"))?;
32
33            match schema.get(&column) {
34                Some(DataType::List(_)) => {
35                    df = df.explode(vec![col(column)]);
36                }
37                Some(DataType::Struct(_)) => {
38                    df = df.unnest([column]);
39                }
40                Some(_) => bail!("unnest error: '{column}' is not a list or struct type"),
41                None => bail!("unnest error: unknown column '{column}'"),
42            }
43        }
44
45        ctx.set_df(df)?;
46    } else {
47        bail!("unnest error: missing input dataframe");
48    }
49
50    Ok(())
51}
dply/engine/unnest.rs

dply/engine/
unnest.rs
