dsq_core/ops/
selection_ops.rs

1use polars::prelude::NewChunkedArray;
2
3use crate::error::{Error, Result};
4use crate::Value;
5
6use super::Operation;
7
8/// Select condition operation for filtering (select(condition))
9pub struct SelectConditionOperation {
10    pub condition_ops: Vec<Box<dyn Operation + Send + Sync>>,
11}
12
13impl SelectConditionOperation {
14    #[must_use]
15    pub fn new(condition_ops: Vec<Box<dyn Operation + Send + Sync>>) -> Self {
16        Self { condition_ops }
17    }
18}
19
20impl Operation for SelectConditionOperation {
21    fn apply(&self, value: &Value) -> Result<Value> {
22        // Evaluate the condition
23        let mut condition_value = value.clone();
24        for op in &self.condition_ops {
25            condition_value = op.apply(&condition_value)?;
26        }
27
28        // Handle DataFrame filtering
29        if let Value::DataFrame(df) = value {
30            if let Value::Series(mask_series) = &condition_value {
31                // Convert the series to a boolean mask
32                let mut mask_vec = Vec::new();
33                for i in 0..mask_series.len() {
34                    match mask_series.get(i) {
35                        Ok(val) => {
36                            let is_true = match val {
37                                polars::prelude::AnyValue::Boolean(b) => b,
38                                polars::prelude::AnyValue::Int64(i) => i != 0,
39                                polars::prelude::AnyValue::Float64(f) => f != 0.0,
40                                polars::prelude::AnyValue::String(s) => !s.is_empty(),
41                                _ => false,
42                            };
43                            mask_vec.push(is_true);
44                        }
45                        Err(_) => mask_vec.push(false),
46                    }
47                }
48
49                if mask_vec.len() == df.height() {
50                    let mask_chunked =
51                        polars::prelude::BooleanChunked::from_slice("mask".into(), &mask_vec);
52                    match df.filter(&mask_chunked) {
53                        Ok(filtered_df) => return Ok(Value::DataFrame(filtered_df)),
54                        Err(e) => {
55                            return Err(Error::Operation(
56                                format!("select() failed to filter DataFrame: {e}").into(),
57                            ));
58                        }
59                    }
60                }
61            }
62        }
63
64        // For non-DataFrame values or when condition is not a series, check if condition is truthy
65        let is_truthy = match condition_value {
66            Value::Bool(b) => b,
67            Value::Int(i) if i != 0 => true,
68            Value::Float(f) if f != 0.0 => true,
69            Value::String(s) if !s.is_empty() => true,
70            Value::Array(arr) if !arr.is_empty() => true,
71            Value::Object(obj) if !obj.is_empty() => true,
72            Value::DataFrame(df) if df.height() > 0 => true,
73            Value::Series(series) if !series.is_empty() => true,
74            _ => false,
75        };
76
77        if is_truthy {
78            Ok(value.clone())
79        } else {
80            Ok(Value::Null)
81        }
82    }
83
84    fn description(&self) -> String {
85        "select with condition".to_string()
86    }
87}
dsq_core/ops/selection_ops.rs

dsq_core/ops/
selection_ops.rs
