spark_connect/
tvf.rs

1//! Table-valued functions (TVF) implementation mirroring `pyspark.sql.connect.tvf.TableValuedFunction`.
2//!
3//! Provides table-valued functions that return DataFrames, such as explode, inline,
4//! range, and other TVF operations.
5
6use crate::column::Column;
7use crate::dataframe::DataFrame;
8use crate::expression::Expression;
9use crate::plan::LogicalPlan;
10use crate::session::SparkSession;
11use crate::types::DataType;
12use spark_connect_core::error::Result;
13
14/// Table-valued functions for Spark SQL.
15///
16/// Mirrors `pyspark.sql.connect.tvf.TableValuedFunction`.
17pub struct TableValuedFunction {
18    /// The session this TVF is bound to
19    session: SparkSession,
20}
21
22impl TableValuedFunction {
23    /// Create a new TableValuedFunction instance.
24    pub(crate) fn new(session: SparkSession) -> Self {
25        TableValuedFunction { session }
26    }
27
28    /// Invoke a Python user-defined table function (UDTF), returning its output
29    /// as a DataFrame. Mirrors `pyspark.sql.functions.udtf`. The Python side
30    /// cloudpickles the handler and sets `eval_type` (300 SQL_TABLE_UDF,
31    /// 301 SQL_ARROW_TABLE_UDF, 302 SQL_ARROW_UDTF); `command` is the pickled payload.
32    #[allow(clippy::too_many_arguments)]
33    pub fn udtf(
34        &self,
35        name: &str,
36        arguments: Vec<Column>,
37        return_type: Option<DataType>,
38        eval_type: i32,
39        command: Vec<u8>,
40        python_ver: String,
41        deterministic: bool,
42    ) -> DataFrame {
43        let args: Vec<Expression> = arguments.iter().map(|c| c.expression().clone()).collect();
44        let plan = LogicalPlan::CommonInlineUdtf {
45            function_name: name.to_string(),
46            deterministic,
47            arguments: args,
48            return_type,
49            eval_type,
50            command,
51            python_ver,
52        };
53        DataFrame::new(self.session.clone(), plan)
54    }
55
56    /// Create a DataFrame representing a range of integers.
57    ///
58    /// This is a table-valued function that generates a sequence of integers
59    /// from `start` to `end` with the given `step`.
60    ///
61    /// # Arguments
62    ///
63    /// * `start` - Starting value (inclusive)
64    /// * `end` - Ending value (exclusive) or the only value if called with one argument
65    /// * `step` - Step size between values (default 1)
66    /// * `num_partitions` - Number of partitions (optional)
67    ///
68    /// # Example
69    ///
70    /// ```ignore
71    /// let df = spark.tvf().range(0, 100, 1, None)?;
72    /// ```
73    pub fn range(
74        &self,
75        start: i64,
76        end: Option<i64>,
77        step: i64,
78        num_partitions: Option<i32>,
79    ) -> Result<DataFrame> {
80        // Handle single-argument case: range(end)
81        let (actual_start, actual_end) = if let Some(e) = end {
82            (start, e)
83        } else {
84            (0, start)
85        };
86
87        self.session
88            .range_full(actual_start, actual_end, step, num_partitions)
89    }
90
91    /// Explode an array or map column.
92    ///
93    /// Returns a DataFrame where each array/map element becomes a separate row.
94    ///
95    /// # Arguments
96    ///
97    /// * `collection` - The column to explode (must be an array or map)
98    ///
99    /// # Example
100    ///
101    /// ```ignore
102    /// let df = spark.tvf().explode(col("my_array"))?;
103    /// ```
104    pub fn explode(&self, collection: &Column) -> Result<DataFrame> {
105        self._fn("explode", vec![collection.clone()])
106    }
107
108    /// Explode an array or map column, preserving nulls as empty rows.
109    ///
110    /// Similar to `explode`, but null values are preserved.
111    pub fn explode_outer(&self, collection: &Column) -> Result<DataFrame> {
112        self._fn("explode_outer", vec![collection.clone()])
113    }
114
115    /// Inline an array of structs column.
116    ///
117    /// Each struct becomes a separate row with its fields as columns.
118    ///
119    /// # Arguments
120    ///
121    /// * `input` - The column to inline (must be an array of structs)
122    ///
123    /// # Example
124    ///
125    /// ```ignore
126    /// let df = spark.tvf().inline(col("struct_array"))?;
127    /// ```
128    pub fn inline(&self, input: &Column) -> Result<DataFrame> {
129        self._fn("inline", vec![input.clone()])
130    }
131
132    /// Inline an array of structs column, preserving nulls as empty rows.
133    ///
134    /// Similar to `inline`, but null values are preserved.
135    pub fn inline_outer(&self, input: &Column) -> Result<DataFrame> {
136        self._fn("inline_outer", vec![input.clone()])
137    }
138
139    /// JSON tuple function.
140    ///
141    /// Extracts fields from a JSON string.
142    ///
143    /// # Arguments
144    ///
145    /// * `input` - The JSON column
146    /// * `fields` - Field names to extract
147    ///
148    /// # Example
149    ///
150    /// ```ignore
151    /// let df = spark.tvf().json_tuple(col("json_str"), vec![col("field1"), col("field2")])?;
152    /// ```
153    pub fn json_tuple<C: Into<Column>>(
154        &self,
155        input: &Column,
156        fields: impl IntoIterator<Item = C>,
157    ) -> Result<DataFrame> {
158        let fields: Vec<Column> = fields.into_iter().map(Into::into).collect();
159        if fields.is_empty() {
160            return Err(spark_connect_core::error::SparkError::value(
161                "CANNOT_BE_EMPTY",
162                &[("item", "field")],
163            ));
164        }
165        let mut args = vec![input.clone()];
166        args.extend(fields);
167        self._fn("json_tuple", args)
168    }
169
170    /// Explode an array column with position indices.
171    ///
172    /// Similar to `explode`, but also includes the position of each element.
173    pub fn posexplode(&self, collection: &Column) -> Result<DataFrame> {
174        self._fn("posexplode", vec![collection.clone()])
175    }
176
177    /// Explode an array column with position indices, preserving nulls.
178    ///
179    /// Similar to `posexplode`, but null values are preserved.
180    pub fn posexplode_outer(&self, collection: &Column) -> Result<DataFrame> {
181        self._fn("posexplode_outer", vec![collection.clone()])
182    }
183
184    /// Stack function.
185    ///
186    /// Converts multiple columns into rows.
187    ///
188    /// # Arguments
189    ///
190    /// * `n` - The number of columns to stack
191    /// * `fields` - The columns to stack
192    pub fn stack<C: Into<Column>>(
193        &self,
194        n: &Column,
195        fields: impl IntoIterator<Item = C>,
196    ) -> Result<DataFrame> {
197        let fields: Vec<Column> = fields.into_iter().map(Into::into).collect();
198        let mut args = vec![n.clone()];
199        args.extend(fields);
200        self._fn("stack", args)
201    }
202
203    /// Returns a DataFrame of collation names.
204    pub fn collations(&self) -> Result<DataFrame> {
205        self._fn("collations", vec![])
206    }
207
208    /// Returns a DataFrame of SQL keywords.
209    pub fn sql_keywords(&self) -> Result<DataFrame> {
210        self._fn("sql_keywords", vec![])
211    }
212
213    /// Explode a variant column.
214    ///
215    /// Variant columns contain semi-structured data in Spark SQL.
216    pub fn variant_explode(&self, input: &Column) -> Result<DataFrame> {
217        self._fn("variant_explode", vec![input.clone()])
218    }
219
220    /// Explode a variant column, preserving nulls.
221    pub fn variant_explode_outer(&self, input: &Column) -> Result<DataFrame> {
222        self._fn("variant_explode_outer", vec![input.clone()])
223    }
224
225    /// Get Python worker logs as a DataFrame.
226    pub fn python_worker_logs(&self) -> Result<DataFrame> {
227        self._fn("python_worker_logs", vec![])
228    }
229
230    /// Generic function implementation.
231    ///
232    /// Internal helper to create an UnresolvedTableValuedFunction.
233    fn _fn(&self, name: &str, args: Vec<Column>) -> Result<DataFrame> {
234        let expressions: Vec<Expression> = args.iter().map(|c| c.expression().clone()).collect();
235        let plan = LogicalPlan::UnresolvedTableValuedFunction {
236            name: name.to_string(),
237            arguments: expressions,
238        };
239        Ok(DataFrame::new(self.session.clone(), plan))
240    }
241}
242
243#[cfg(test)]
244mod tests {
245    use super::*;
246    use crate::session::SparkSession;
247
248    fn session() -> SparkSession {
249        SparkSession::builder()
250            .remote("sc://localhost:15002")
251            .get_or_create()
252            .expect("failed to build session")
253    }
254
255    #[test]
256    fn tvf_range_plan() {
257        let spark = session();
258        let tvf = spark.tvf();
259        let df = tvf.range(0, Some(10), 1, None).unwrap();
260        match &df.plan {
261            LogicalPlan::Range { .. } => {
262                // Plan is correct
263            }
264            _ => panic!("expected Range plan"),
265        }
266    }
267
268    #[test]
269    fn tvf_range_single_arg() {
270        let spark = session();
271        let tvf = spark.tvf();
272        // range(end) should be range(0, end)
273        let df = tvf.range(5, None, 1, None).unwrap();
274        match &df.plan {
275            LogicalPlan::Range {
276                start, end, step, ..
277            } => {
278                assert_eq!(*start, 0);
279                assert_eq!(*end, 5);
280                assert_eq!(*step, 1);
281            }
282            _ => panic!("expected Range plan"),
283        }
284    }
285
286    #[test]
287    fn tvf_explode_plan() {
288        let spark = session();
289        let tvf = spark.tvf();
290        let col = crate::column::col("array_col");
291        let df = tvf.explode(&col).unwrap();
292        match &df.plan {
293            LogicalPlan::UnresolvedTableValuedFunction { name, arguments } => {
294                assert_eq!(name, "explode");
295                assert_eq!(arguments.len(), 1);
296            }
297            _ => panic!("expected UnresolvedTableValuedFunction plan"),
298        }
299    }
300
301    #[test]
302    fn tvf_explode_outer_plan() {
303        let spark = session();
304        let tvf = spark.tvf();
305        let col = crate::column::col("array_col");
306        let df = tvf.explode_outer(&col).unwrap();
307        match &df.plan {
308            LogicalPlan::UnresolvedTableValuedFunction { name, .. } => {
309                assert_eq!(name, "explode_outer");
310            }
311            _ => panic!("expected UnresolvedTableValuedFunction plan"),
312        }
313    }
314
315    #[test]
316    fn tvf_inline_plan() {
317        let spark = session();
318        let tvf = spark.tvf();
319        let col = crate::column::col("struct_array");
320        let df = tvf.inline(&col).unwrap();
321        match &df.plan {
322            LogicalPlan::UnresolvedTableValuedFunction { name, .. } => {
323                assert_eq!(name, "inline");
324            }
325            _ => panic!("expected UnresolvedTableValuedFunction plan"),
326        }
327    }
328
329    #[test]
330    fn tvf_json_tuple_plan() {
331        let spark = session();
332        let tvf = spark.tvf();
333        let input = crate::column::col("json_col");
334        let fields = vec![crate::column::col("field1"), crate::column::col("field2")];
335        let df = tvf.json_tuple(&input, fields).unwrap();
336        match &df.plan {
337            LogicalPlan::UnresolvedTableValuedFunction { name, arguments } => {
338                assert_eq!(name, "json_tuple");
339                assert_eq!(arguments.len(), 3); // input + 2 fields
340            }
341            _ => panic!("expected UnresolvedTableValuedFunction plan"),
342        }
343    }
344
345    #[test]
346    fn tvf_json_tuple_empty_fields_error() {
347        let spark = session();
348        let tvf = spark.tvf();
349        let input = crate::column::col("json_col");
350        let result = tvf.json_tuple(&input, Vec::<Column>::new());
351        assert!(result.is_err());
352    }
353
354    #[test]
355    fn tvf_posexplode_plan() {
356        let spark = session();
357        let tvf = spark.tvf();
358        let col = crate::column::col("array_col");
359        let df = tvf.posexplode(&col).unwrap();
360        match &df.plan {
361            LogicalPlan::UnresolvedTableValuedFunction { name, .. } => {
362                assert_eq!(name, "posexplode");
363            }
364            _ => panic!("expected UnresolvedTableValuedFunction plan"),
365        }
366    }
367
368    #[test]
369    fn tvf_stack_plan() {
370        let spark = session();
371        let tvf = spark.tvf();
372        let n = crate::column::lit(3i64);
373        let fields = vec![
374            crate::column::col("col1"),
375            crate::column::col("col2"),
376            crate::column::col("col3"),
377        ];
378        let df = tvf.stack(&n, fields).unwrap();
379        match &df.plan {
380            LogicalPlan::UnresolvedTableValuedFunction { name, arguments } => {
381                assert_eq!(name, "stack");
382                assert_eq!(arguments.len(), 4); // n + 3 fields
383            }
384            _ => panic!("expected UnresolvedTableValuedFunction plan"),
385        }
386    }
387
388    #[test]
389    fn tvf_collations_plan() {
390        let spark = session();
391        let tvf = spark.tvf();
392        let df = tvf.collations().unwrap();
393        match &df.plan {
394            LogicalPlan::UnresolvedTableValuedFunction { name, arguments } => {
395                assert_eq!(name, "collations");
396                assert!(arguments.is_empty());
397            }
398            _ => panic!("expected UnresolvedTableValuedFunction plan"),
399        }
400    }
401
402    #[test]
403    fn tvf_sql_keywords_plan() {
404        let spark = session();
405        let tvf = spark.tvf();
406        let df = tvf.sql_keywords().unwrap();
407        match &df.plan {
408            LogicalPlan::UnresolvedTableValuedFunction { name, arguments } => {
409                assert_eq!(name, "sql_keywords");
410                assert!(arguments.is_empty());
411            }
412            _ => panic!("expected UnresolvedTableValuedFunction plan"),
413        }
414    }
415}
spark_connect/tvf.rs

spark_connect/
tvf.rs
