/home/noah/src/ruchy/src/backend/transpiler/dataframe.rs

Source
//! `DataFrame` transpilation for Polars integration

#![allow(clippy::missing_errors_doc)]
#![allow(clippy::wildcard_imports)]
#![allow(clippy::doc_markdown)]

use super::*;
use crate::frontend::ast::{AggregateOp, DataFrameColumn, DataFrameOp, JoinType};
use anyhow::Result;
use proc_macro2::TokenStream;
use quote::{format_ident, quote};

impl Transpiler {
    /// Transpiles DataFrame literals (df![] syntax)
    pub fn transpile_dataframe(&self, columns: &[DataFrameColumn]) -> Result<TokenStream> {
        if columns.is_empty() {
            // Empty DataFrame
            return Ok(quote! {
                polars::prelude::DataFrame::empty()
            });
        }

        let mut series_tokens = Vec::new();

        for column in columns {
            let col_name = &column.name;

            // Transpile the column values
            let values_tokens = if column.values.is_empty() {
                quote! { vec![] }
            } else {
                // Collect all values into a vector
                let value_tokens: Result<Vec<_>> = column
                    .values
                    .iter()
                    .map(|v| self.transpile_expr(v))
                    .collect();
                let value_tokens = value_tokens?;
                quote! { vec![#(#value_tokens),*] }
            };

            // Create a Series from the values
            series_tokens.push(quote! {
                polars::prelude::Series::new(#col_name, #values_tokens)
            });
        }

        // Create DataFrame from series
        Ok(quote! {
            polars::prelude::DataFrame::new(vec![
                #(#series_tokens),*
            ]).unwrap()
        })
    }

    /// Transpiles DataFrame operations
    pub fn transpile_dataframe_operation(
        &self,
        df: &Expr,
        op: &DataFrameOp,
    ) -> Result<TokenStream> {
        let df_tokens = self.transpile_expr(df)?;

        match op {
            DataFrameOp::Select(columns) => {
                let col_tokens: Vec<TokenStream> =
                    columns.iter().map(|col| quote! { #col }).collect();
                Ok(quote! {
                    #df_tokens.select(&[#(#col_tokens),*]).unwrap()
                })
            }
            DataFrameOp::Filter(condition) => {
                let cond_tokens = self.transpile_expr(condition)?;
                Ok(quote! {
                    #df_tokens.filter(&#cond_tokens).unwrap()
                })
            }
            DataFrameOp::GroupBy(columns) => {
                let col_tokens: Vec<TokenStream> =
                    columns.iter().map(|col| quote! { #col }).collect();
                Ok(quote! {
                    #df_tokens.groupby(&[#(#col_tokens),*]).unwrap()
                })
            }
            DataFrameOp::Sort(columns) => {
                // Sort by multiple columns
                let col_tokens: Vec<TokenStream> =
                    columns.iter().map(|col| quote! { #col }).collect();
                Ok(quote! {
                    #df_tokens.sort(&[#(#col_tokens),*], false).unwrap()
                })
            }
            DataFrameOp::Join { other, on, how } => {
                let other_tokens = self.transpile_expr(other)?;
                let on_tokens: Vec<TokenStream> = on.iter().map(|col| quote! { #col }).collect();

                let join_type = match how {
                    JoinType::Left => quote! { polars::prelude::JoinType::Left },
                    JoinType::Right => quote! { polars::prelude::JoinType::Right },
                    JoinType::Inner => quote! { polars::prelude::JoinType::Inner },
                    JoinType::Outer => quote! { polars::prelude::JoinType::Outer },
                };

                Ok(quote! {
                    #df_tokens.join(
                        &#other_tokens,
                        &[#(#on_tokens),*],
                        &[#(#on_tokens),*],
                        #join_type
                    ).unwrap()
                })
            }
            DataFrameOp::Aggregate(agg_ops) => {
                // Convert AggregateOp to expressions
                let agg_exprs: Vec<TokenStream> = agg_ops
                    .iter()
                    .map(|op| match op {
                        AggregateOp::Sum(col) => quote! { col(#col).sum() },
                        AggregateOp::Mean(col) => quote! { col(#col).mean() },
                        AggregateOp::Min(col) => quote! { col(#col).min() },
                        AggregateOp::Max(col) => quote! { col(#col).max() },
                        AggregateOp::Count(col) => quote! { col(#col).count() },
                        AggregateOp::Std(col) => quote! { col(#col).std() },
                        AggregateOp::Var(col) => quote! { col(#col).var() },
                    })
                    .collect();

                Ok(quote! {
                    #df_tokens.agg(&[#(#agg_exprs),*]).unwrap()
                })
            }
            DataFrameOp::Limit(n) => Ok(quote! {
                #df_tokens.limit(#n)
            }),
            DataFrameOp::Head(n) => Ok(quote! {
                #df_tokens.head(Some(#n))
            }),
            DataFrameOp::Tail(n) => Ok(quote! {
                #df_tokens.tail(Some(#n))
            }),
        }
    }

    /// Transpiles DataFrame method calls (alternative to operation enum)
    pub fn transpile_dataframe_method(
        &self,
        df_expr: &Expr,
        method: &str,
        args: &[Expr],
    ) -> Result<TokenStream> {
        let df_tokens = self.transpile_expr(df_expr)?;

        let arg_tokens: Result<Vec<_>> = args.iter().map(|a| self.transpile_expr(a)).collect();
        let arg_tokens = arg_tokens?;

        // Map Ruchy DataFrame methods to correct Polars API
        match method {
            // DataFrame builder pattern methods (don't exist in Polars)
            "column" => {
                // .column() doesn't exist - this is part of builder pattern
                // For now, just pass through to show the issue
                if arg_tokens.len() == 2 {
                    let name = &arg_tokens[0];
                    let data = &arg_tokens[1];
                    Ok(quote! { #df_tokens.column(#name, #data) })
                } else {
                    Ok(quote! { #df_tokens.column(#(#arg_tokens),*) })
                }
            }
            "build" => {
                // .build() doesn't exist in Polars - just return the DataFrame
                Ok(quote! { #df_tokens })
            }
            
            // DataFrame inspection methods
            "rows" => {
                // Polars uses .height() not .rows()
                Ok(quote! { #df_tokens.height() })
            }
            "columns" => {
                // Polars uses .get_column_names() for column names
                Ok(quote! { #df_tokens.get_column_names() })
            }
            "get" => {
                // df.get("column") -> df.column("column") in Polars
                if arg_tokens.len() == 1 {
                    let col_name = &arg_tokens[0];
                    Ok(quote! { #df_tokens.column(#col_name) })
                } else {
                    Ok(quote! { #df_tokens.get(#(#arg_tokens),*) })
                }
            }
            
            // DataFrame operations
            "select" | "filter" | "sort" => {
                // These need lazy evaluation for proper chaining
                let method_ident = format_ident!("{}", method);
                Ok(quote! {
                    #df_tokens.lazy().#method_ident(#(#arg_tokens),*).collect().unwrap()
                })
            }
            "groupby" | "group_by" => {
                // Polars uses group_by
                Ok(quote! {
                    #df_tokens.group_by(#(#arg_tokens),*).unwrap()
                })
            }
            "agg" | "join" => {
                let method_ident = format_ident!("{}", method);
                Ok(quote! {
                    #df_tokens.#method_ident(#(#arg_tokens),*).unwrap()
                })
            }
            
            // Statistical methods
            "mean" | "std" | "min" | "max" | "sum" | "count" => {
                // These are aggregate functions
                let method_ident = format_ident!("{}", method);
                Ok(quote! {
                    #df_tokens.#method_ident()
                })
            }
            
            // Head/tail methods
            "head" | "tail" => {
                let method_ident = format_ident!("{}", method);
                if args.is_empty() {
                    Ok(quote! { #df_tokens.#method_ident(Some(5)) })
                } else {
                    Ok(quote! { #df_tokens.#method_ident(Some(#(#arg_tokens),*)) })
                }
            }
            _ => {
                // Default method call
                let method_ident = format_ident!("{}", method);
                Ok(quote! {
                    #df_tokens.#method_ident(#(#arg_tokens),*)
                })
            }
        }
    }
    
    /// Check if an expression is a DataFrame type
    pub fn is_dataframe_expr(&self, expr: &Expr) -> bool {
        use crate::frontend::ast::ExprKind;
        
        match &expr.kind {
            // Variable named "df" is likely a DataFrame
            ExprKind::Identifier(name) if name == "df" => true,
            
            // DataFrame constructor calls
            ExprKind::Call { func, .. } => {
                if let ExprKind::QualifiedName { module, name } = &func.kind {
                    module == "DataFrame" && (name == "new" || name == "from_csv" || name == "from_json" || name == "from_csv_string")
                } else {
                    false
                }
            }
            
            // Method calls that return DataFrames
            ExprKind::MethodCall { receiver, method, .. } => {
                // Check if it's a DataFrame method chain
                matches!(method.as_str(), 
                    "column" | "build" | "select" | "filter" | "sort" | 
                    "head" | "tail" | "drop_nulls" | "fill_null"
                ) || self.is_dataframe_expr(receiver)
            }
            
            // DataFrame literals
            ExprKind::DataFrame { .. } => true,
            
            _ => false,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::frontend::ast::{Expr, ExprKind, Literal, Span};
    
    fn make_test_transpiler() -> Transpiler {
        Transpiler::new()
    }
    
    fn make_literal_expr(val: i64) -> Expr {
        Expr {
            kind: ExprKind::Literal(Literal::Integer(val)),
            span: Span::new(0, 10),
            attributes: vec![],
        }
    }
    
    #[test]
    fn test_empty_dataframe() {
        let transpiler = make_test_transpiler();
        let result = transpiler.transpile_dataframe(&[]).unwrap();
        let output = result.to_string();
        assert!(output.contains("DataFrame"));
        assert!(output.contains("empty"));
    }
    
    #[test]
    fn test_dataframe_with_columns() {
        let transpiler = make_test_transpiler();
        let columns = vec![
            DataFrameColumn {
                name: "col1".to_string(),
                values: vec![make_literal_expr(1), make_literal_expr(2)],
            },
            DataFrameColumn {
                name: "col2".to_string(),
                values: vec![make_literal_expr(3), make_literal_expr(4)],
            },
        ];
        
        let result = transpiler.transpile_dataframe(&columns).unwrap();
        let output = result.to_string();
        assert!(output.contains("DataFrame"));
        assert!(output.contains("Series"));
        assert!(output.contains("col1"));
        assert!(output.contains("col2"));
    }
    
    #[test]
    fn test_dataframe_select_operation() {
        let transpiler = make_test_transpiler();
        let df_expr = make_literal_expr(0); // Placeholder
        let op = DataFrameOp::Select(vec!["col1".to_string(), "col2".to_string()]);
        
        let result = transpiler.transpile_dataframe_operation(&df_expr, &op).unwrap();
        let output = result.to_string();
        assert!(output.contains("select"));
        assert!(output.contains("col1"));
        assert!(output.contains("col2"));
    }
    
    #[test]
    fn test_dataframe_filter_operation() {
        let transpiler = make_test_transpiler();
        let df_expr = make_literal_expr(0);
        let condition = make_literal_expr(1);
        let op = DataFrameOp::Filter(Box::new(condition));
        
        let result = transpiler.transpile_dataframe_operation(&df_expr, &op).unwrap();
        let output = result.to_string();
        assert!(output.contains("filter"));
    }
    
    #[test]
    fn test_dataframe_groupby_operation() {
        let transpiler = make_test_transpiler();
        let df_expr = make_literal_expr(0);
        let op = DataFrameOp::GroupBy(vec!["group_col".to_string()]);
        
        let result = transpiler.transpile_dataframe_operation(&df_expr, &op).unwrap();
        let output = result.to_string();
        assert!(output.contains("groupby"));
        assert!(output.contains("group_col"));
    }
    
    #[test]
    fn test_dataframe_sort_operation() {
        let transpiler = make_test_transpiler();
        let df_expr = make_literal_expr(0);
        let op = DataFrameOp::Sort(vec!["sort_col".to_string()]);
        
        let result = transpiler.transpile_dataframe_operation(&df_expr, &op).unwrap();
        let output = result.to_string();
        assert!(output.contains("sort"));
        assert!(output.contains("sort_col"));
    }
    
    #[test]
    fn test_dataframe_join_operations() {
        let transpiler = make_test_transpiler();
        let df_expr = make_literal_expr(0);
        let other_expr = make_literal_expr(1);
        
        let join_types = vec![
            (JoinType::Inner, "Inner"),
            (JoinType::Left, "Left"),
            (JoinType::Right, "Right"),
        ];
        
        for (join_type, expected) in join_types {
            let op = DataFrameOp::Join {
                other: Box::new(other_expr.clone()),
                on: vec!["id".to_string()],
                how: join_type,
            };
            
            let result = transpiler.transpile_dataframe_operation(&df_expr, &op).unwrap();
            let output = result.to_string();
            assert!(output.contains("join"));
            assert!(output.contains(expected));
        }
    }
    
    #[test]
    fn test_dataframe_aggregate_operations() {
        let transpiler = make_test_transpiler();
        let df_expr = make_literal_expr(0);
        
        let agg_ops = vec![
            AggregateOp::Mean("col1".to_string()),
            AggregateOp::Sum("col2".to_string()),
            AggregateOp::Min("col3".to_string()),
            AggregateOp::Max("col4".to_string()),
            AggregateOp::Count("col5".to_string()),
            AggregateOp::Std("col6".to_string()),
        ];
        
        let op = DataFrameOp::Aggregate(agg_ops);
        let result = transpiler.transpile_dataframe_operation(&df_expr, &op).unwrap();
        let output = result.to_string();
        // Check that it produces some output
        assert!(!output.is_empty());
    }
    
    #[test]
    fn test_dataframe_limit_operations() {
        let transpiler = make_test_transpiler();
        let df_expr = make_literal_expr(0);
        
        // Test Limit
        let op = DataFrameOp::Limit(10);
        let result = transpiler.transpile_dataframe_operation(&df_expr, &op).unwrap();
        let output = result.to_string();
        assert!(output.contains("limit"));
        
        // Test Head
        let op = DataFrameOp::Head(5);
        let result = transpiler.transpile_dataframe_operation(&df_expr, &op).unwrap();
        let output = result.to_string();
        assert!(output.contains("head"));
        
        // Test Tail
        let op = DataFrameOp::Tail(5);
        let result = transpiler.transpile_dataframe_operation(&df_expr, &op).unwrap();
        let output = result.to_string();
        assert!(output.contains("tail"));
    }
    
    #[test]
    fn test_dataframe_with_empty_column_values() {
        let transpiler = make_test_transpiler();
        let columns = vec![
            DataFrameColumn {
                name: "empty_col".to_string(),
                values: vec![],
            },
        ];
        
        let result = transpiler.transpile_dataframe(&columns).unwrap();
        let output = result.to_string();
        assert!(output.contains("Series"));
        assert!(output.contains("empty_col"));
        assert!(output.contains("vec"));
    }
}

Coverage Report

Created: 2025-09-08 21:26

Line	Count	Source
1		//! `DataFrame` transpilation for Polars integration
2
3		#![allow(clippy::missing_errors_doc)]
4		#![allow(clippy::wildcard_imports)]
5		#![allow(clippy::doc_markdown)]
6
7		use super::*;
8		use crate::frontend::ast::{AggregateOp, DataFrameColumn, DataFrameOp, JoinType};
9		use anyhow::Result;
10		use proc_macro2::TokenStream;
11		use quote::{format_ident, quote};
12
13		impl Transpiler {
14		/// Transpiles DataFrame literals (df![] syntax)
15	0	pub fn transpile_dataframe(&self, columns: &[DataFrameColumn]) -> Result<TokenStream> {
16	0	if columns.is_empty() {
17		// Empty DataFrame
18	0	return Ok(quote! {
19	0	polars::prelude::DataFrame::empty()
20	0	});
21	0	}
22
23	0	let mut series_tokens = Vec::new();
24
25	0	for column in columns {
26	0	let col_name = &column.name;
27
28		// Transpile the column values
29	0	let values_tokens = if column.values.is_empty() {
30	0	quote! { vec![] }
31		} else {
32		// Collect all values into a vector
33	0	let value_tokens: Result<Vec<_>> = column
34	0	.values
35	0	.iter()
36	0	.map(\|v\| self.transpile_expr(v))
37	0	.collect();
38	0	let value_tokens = value_tokens?;
39	0	quote! { vec![#(#value_tokens),*] }
40		};
41
42		// Create a Series from the values
43	0	series_tokens.push(quote! {
44		polars::prelude::Series::new(#col_name, #values_tokens)
45		});
46		}
47
48		// Create DataFrame from series
49	0	Ok(quote! {
50		polars::prelude::DataFrame::new(vec![
51		#(#series_tokens),*
52		]).unwrap()
53		})
54	0	}
55
56		/// Transpiles DataFrame operations
57	0	pub fn transpile_dataframe_operation(
58	0	&self,
59	0	df: &Expr,
60	0	op: &DataFrameOp,
61	0	) -> Result<TokenStream> {
62	0	let df_tokens = self.transpile_expr(df)?;
63
64	0	match op {
65	0	DataFrameOp::Select(columns) => {
66	0	let col_tokens: Vec<TokenStream> =
67	0	columns.iter().map(\|col\| quote! { #col }).collect();
68	0	Ok(quote! {
69		#df_tokens.select(&[#(#col_tokens),*]).unwrap()
70		})
71		}
72	0	DataFrameOp::Filter(condition) => {
73	0	let cond_tokens = self.transpile_expr(condition)?;
74	0	Ok(quote! {
75	0	#df_tokens.filter(&#cond_tokens).unwrap()
76	0	})
77		}
78	0	DataFrameOp::GroupBy(columns) => {
79	0	let col_tokens: Vec<TokenStream> =
80	0	columns.iter().map(\|col\| quote! { #col }).collect();
81	0	Ok(quote! {
82		#df_tokens.groupby(&[#(#col_tokens),*]).unwrap()
83		})
84		}
85	0	DataFrameOp::Sort(columns) => {
86		// Sort by multiple columns
87	0	let col_tokens: Vec<TokenStream> =
88	0	columns.iter().map(\|col\| quote! { #col }).collect();
89	0	Ok(quote! {
90		#df_tokens.sort(&[#(#col_tokens),*], false).unwrap()
91		})
92		}
93	0	DataFrameOp::Join { other, on, how } => {
94	0	let other_tokens = self.transpile_expr(other)?;
95	0	let on_tokens: Vec<TokenStream> = on.iter().map(\|col\| quote! { #col }).collect();
96
97	0	let join_type = match how {
98	0	JoinType::Left => quote! { polars::prelude::JoinType::Left },
99	0	JoinType::Right => quote! { polars::prelude::JoinType::Right },
100	0	JoinType::Inner => quote! { polars::prelude::JoinType::Inner },
101	0	JoinType::Outer => quote! { polars::prelude::JoinType::Outer },
102		};
103
104	0	Ok(quote! {
105		#df_tokens.join(
106		&#other_tokens,
107		&[#(#on_tokens),*],
108		&[#(#on_tokens),*],
109		#join_type
110		).unwrap()
111		})
112		}
113	0	DataFrameOp::Aggregate(agg_ops) => {
114		// Convert AggregateOp to expressions
115	0	let agg_exprs: Vec<TokenStream> = agg_ops
116	0	.iter()
117	0	.map(\|op\| match op {
118	0	AggregateOp::Sum(col) => quote! { col(#col).sum() },
119	0	AggregateOp::Mean(col) => quote! { col(#col).mean() },
120	0	AggregateOp::Min(col) => quote! { col(#col).min() },
121	0	AggregateOp::Max(col) => quote! { col(#col).max() },
122	0	AggregateOp::Count(col) => quote! { col(#col).count() },
123	0	AggregateOp::Std(col) => quote! { col(#col).std() },
124	0	AggregateOp::Var(col) => quote! { col(#col).var() },
125	0	})
126	0	.collect();
127
128	0	Ok(quote! {
129		#df_tokens.agg(&[#(#agg_exprs),*]).unwrap()
130		})
131		}
132	0	DataFrameOp::Limit(n) => Ok(quote! {
133	0	#df_tokens.limit(#n)
134	0	}),
135	0	DataFrameOp::Head(n) => Ok(quote! {
136	0	#df_tokens.head(Some(#n))
137	0	}),
138	0	DataFrameOp::Tail(n) => Ok(quote! {
139	0	#df_tokens.tail(Some(#n))
140	0	}),
141		}
142	0	}
143
144		/// Transpiles DataFrame method calls (alternative to operation enum)
145	0	pub fn transpile_dataframe_method(
146	0	&self,
147	0	df_expr: &Expr,
148	0	method: &str,
149	0	args: &[Expr],
150	0	) -> Result<TokenStream> {
151	0	let df_tokens = self.transpile_expr(df_expr)?;
152
153	0	let arg_tokens: Result<Vec<_>> = args.iter().map(\|a\| self.transpile_expr(a)).collect();
154	0	let arg_tokens = arg_tokens?;
155
156		// Map Ruchy DataFrame methods to correct Polars API
157	0	match method {
158		// DataFrame builder pattern methods (don't exist in Polars)
159	0	"column" => {
160		// .column() doesn't exist - this is part of builder pattern
161		// For now, just pass through to show the issue
162	0	if arg_tokens.len() == 2 {
163	0	let name = &arg_tokens[0];
164	0	let data = &arg_tokens[1];
165	0	Ok(quote! { #df_tokens.column(#name, #data) })
166		} else {
167	0	Ok(quote! { #df_tokens.column(#(#arg_tokens),*) })
168		}
169		}
170	0	"build" => {
171		// .build() doesn't exist in Polars - just return the DataFrame
172	0	Ok(quote! { #df_tokens })
173		}
174
175		// DataFrame inspection methods
176	0	"rows" => {
177		// Polars uses .height() not .rows()
178	0	Ok(quote! { #df_tokens.height() })
179		}
180	0	"columns" => {
181		// Polars uses .get_column_names() for column names
182	0	Ok(quote! { #df_tokens.get_column_names() })
183		}
184	0	"get" => {
185		// df.get("column") -> df.column("column") in Polars
186	0	if arg_tokens.len() == 1 {
187	0	let col_name = &arg_tokens[0];
188	0	Ok(quote! { #df_tokens.column(#col_name) })
189		} else {
190	0	Ok(quote! { #df_tokens.get(#(#arg_tokens),*) })
191		}
192		}
193
194		// DataFrame operations
195	0	"select" \| "filter" \| "sort" => {
196		// These need lazy evaluation for proper chaining
197	0	let method_ident = format_ident!("{}", method);
198	0	Ok(quote! {
199		#df_tokens.lazy().#method_ident(#(#arg_tokens),*).collect().unwrap()
200		})
201		}
202	0	"groupby" \| "group_by" => {
203		// Polars uses group_by
204	0	Ok(quote! {
205		#df_tokens.group_by(#(#arg_tokens),*).unwrap()
206		})
207		}
208	0	"agg" \| "join" => {
209	0	let method_ident = format_ident!("{}", method);
210	0	Ok(quote! {
211		#df_tokens.#method_ident(#(#arg_tokens),*).unwrap()
212		})
213		}
214
215		// Statistical methods
216	0	"mean" \| "std" \| "min" \| "max" \| "sum" \| "count" => {
217		// These are aggregate functions
218	0	let method_ident = format_ident!("{}", method);
219	0	Ok(quote! {
220	0	#df_tokens.#method_ident()
221	0	})
222		}
223
224		// Head/tail methods
225	0	"head" \| "tail" => {
226	0	let method_ident = format_ident!("{}", method);
227	0	if args.is_empty() {
228	0	Ok(quote! { #df_tokens.#method_ident(Some(5)) })
229		} else {
230	0	Ok(quote! { #df_tokens.#method_ident(Some(#(#arg_tokens),*)) })
231		}
232		}
233		_ => {
234		// Default method call
235	0	let method_ident = format_ident!("{}", method);
236	0	Ok(quote! {
237		#df_tokens.#method_ident(#(#arg_tokens),*)
238		})
239		}
240		}
241	0	}
242
243		/// Check if an expression is a DataFrame type
244	0	pub fn is_dataframe_expr(&self, expr: &Expr) -> bool {
245		use crate::frontend::ast::ExprKind;
246
247	0	match &expr.kind {
248		// Variable named "df" is likely a DataFrame
249	0	ExprKind::Identifier(name) if name == "df" => true,
250
251		// DataFrame constructor calls
252	0	ExprKind::Call { func, .. } => {
253	0	if let ExprKind::QualifiedName { module, name } = &func.kind {
254	0	module == "DataFrame" && (name == "new" \|\| name == "from_csv" \|\| name == "from_json" \|\| name == "from_csv_string")
255		} else {
256	0	false
257		}
258		}
259
260		// Method calls that return DataFrames
261	0	ExprKind::MethodCall { receiver, method, .. } => {
262		// Check if it's a DataFrame method chain
263	0	matches!(method.as_str(),
264	0	"column" \| "build" \| "select" \| "filter" \| "sort" \|
265	0	"head" \| "tail" \| "drop_nulls" \| "fill_null"
266	0	) \|\| self.is_dataframe_expr(receiver)
267		}
268
269		// DataFrame literals
270	0	ExprKind::DataFrame { .. } => true,
271
272	0	_ => false,
273		}
274	0	}
275		}
276
277		#[cfg(test)]
278		mod tests {
279		use super::*;
280		use crate::frontend::ast::{Expr, ExprKind, Literal, Span};
281
282		fn make_test_transpiler() -> Transpiler {
283		Transpiler::new()
284		}
285
286		fn make_literal_expr(val: i64) -> Expr {
287		Expr {
288		kind: ExprKind::Literal(Literal::Integer(val)),
289		span: Span::new(0, 10),
290		attributes: vec![],
291		}
292		}
293
294		#[test]
295		fn test_empty_dataframe() {
296		let transpiler = make_test_transpiler();
297		let result = transpiler.transpile_dataframe(&[]).unwrap();
298		let output = result.to_string();
299		assert!(output.contains("DataFrame"));
300		assert!(output.contains("empty"));
301		}
302
303		#[test]
304		fn test_dataframe_with_columns() {
305		let transpiler = make_test_transpiler();
306		let columns = vec![
307		DataFrameColumn {
308		name: "col1".to_string(),
309		values: vec![make_literal_expr(1), make_literal_expr(2)],
310		},
311		DataFrameColumn {
312		name: "col2".to_string(),
313		values: vec![make_literal_expr(3), make_literal_expr(4)],
314		},
315		];
316
317		let result = transpiler.transpile_dataframe(&columns).unwrap();
318		let output = result.to_string();
319		assert!(output.contains("DataFrame"));
320		assert!(output.contains("Series"));
321		assert!(output.contains("col1"));
322		assert!(output.contains("col2"));
323		}
324
325		#[test]
326		fn test_dataframe_select_operation() {
327		let transpiler = make_test_transpiler();
328		let df_expr = make_literal_expr(0); // Placeholder
329		let op = DataFrameOp::Select(vec!["col1".to_string(), "col2".to_string()]);
330
331		let result = transpiler.transpile_dataframe_operation(&df_expr, &op).unwrap();
332		let output = result.to_string();
333		assert!(output.contains("select"));
334		assert!(output.contains("col1"));
335		assert!(output.contains("col2"));
336		}
337
338		#[test]
339		fn test_dataframe_filter_operation() {
340		let transpiler = make_test_transpiler();
341		let df_expr = make_literal_expr(0);
342		let condition = make_literal_expr(1);
343		let op = DataFrameOp::Filter(Box::new(condition));
344
345		let result = transpiler.transpile_dataframe_operation(&df_expr, &op).unwrap();
346		let output = result.to_string();
347		assert!(output.contains("filter"));
348		}
349
350		#[test]
351		fn test_dataframe_groupby_operation() {
352		let transpiler = make_test_transpiler();
353		let df_expr = make_literal_expr(0);
354		let op = DataFrameOp::GroupBy(vec!["group_col".to_string()]);
355
356		let result = transpiler.transpile_dataframe_operation(&df_expr, &op).unwrap();
357		let output = result.to_string();
358		assert!(output.contains("groupby"));
359		assert!(output.contains("group_col"));
360		}
361
362		#[test]
363		fn test_dataframe_sort_operation() {
364		let transpiler = make_test_transpiler();
365		let df_expr = make_literal_expr(0);
366		let op = DataFrameOp::Sort(vec!["sort_col".to_string()]);
367
368		let result = transpiler.transpile_dataframe_operation(&df_expr, &op).unwrap();
369		let output = result.to_string();
370		assert!(output.contains("sort"));
371		assert!(output.contains("sort_col"));
372		}
373
374		#[test]
375		fn test_dataframe_join_operations() {
376		let transpiler = make_test_transpiler();
377		let df_expr = make_literal_expr(0);
378		let other_expr = make_literal_expr(1);
379
380		let join_types = vec![
381		(JoinType::Inner, "Inner"),
382		(JoinType::Left, "Left"),
383		(JoinType::Right, "Right"),
384		];
385
386		for (join_type, expected) in join_types {
387		let op = DataFrameOp::Join {
388		other: Box::new(other_expr.clone()),
389		on: vec!["id".to_string()],
390		how: join_type,
391		};
392
393		let result = transpiler.transpile_dataframe_operation(&df_expr, &op).unwrap();
394		let output = result.to_string();
395		assert!(output.contains("join"));
396		assert!(output.contains(expected));
397		}
398		}
399
400		#[test]
401		fn test_dataframe_aggregate_operations() {
402		let transpiler = make_test_transpiler();
403		let df_expr = make_literal_expr(0);
404
405		let agg_ops = vec![
406		AggregateOp::Mean("col1".to_string()),
407		AggregateOp::Sum("col2".to_string()),
408		AggregateOp::Min("col3".to_string()),
409		AggregateOp::Max("col4".to_string()),
410		AggregateOp::Count("col5".to_string()),
411		AggregateOp::Std("col6".to_string()),
412		];
413
414		let op = DataFrameOp::Aggregate(agg_ops);
415		let result = transpiler.transpile_dataframe_operation(&df_expr, &op).unwrap();
416		let output = result.to_string();
417		// Check that it produces some output
418		assert!(!output.is_empty());
419		}
420
421		#[test]
422		fn test_dataframe_limit_operations() {
423		let transpiler = make_test_transpiler();
424		let df_expr = make_literal_expr(0);
425
426		// Test Limit
427		let op = DataFrameOp::Limit(10);
428		let result = transpiler.transpile_dataframe_operation(&df_expr, &op).unwrap();
429		let output = result.to_string();
430		assert!(output.contains("limit"));
431
432		// Test Head
433		let op = DataFrameOp::Head(5);
434		let result = transpiler.transpile_dataframe_operation(&df_expr, &op).unwrap();
435		let output = result.to_string();
436		assert!(output.contains("head"));
437
438		// Test Tail
439		let op = DataFrameOp::Tail(5);
440		let result = transpiler.transpile_dataframe_operation(&df_expr, &op).unwrap();
441		let output = result.to_string();
442		assert!(output.contains("tail"));
443		}
444
445		#[test]
446		fn test_dataframe_with_empty_column_values() {
447		let transpiler = make_test_transpiler();
448		let columns = vec![
449		DataFrameColumn {
450		name: "empty_col".to_string(),
451		values: vec![],
452		},
453		];
454
455		let result = transpiler.transpile_dataframe(&columns).unwrap();
456		let output = result.to_string();
457		assert!(output.contains("Series"));
458		assert!(output.contains("empty_col"));
459		assert!(output.contains("vec"));
460		}
461		}