[SPARK-61180][CONNECT][PYTHON] Add DataFrame._joinAsOf, the as-of join of pandas-on-Spark's merge_asof, in the Rust and Python clients as PySpark does

crates/pyspark-rs/src/dataframe.rs

@@ -2,7 +2,7 @@
 
 use pyo3::prelude::*;
 use pyo3::types::{PyDict, PyInt, PyList, PyTuple};
-use spark_connect::dataframe::{DataFrame, LocalRowIterator};
+use spark_connect::dataframe::{AsOfJoinOptions, DataFrame, LocalRowIterator};
 use spark_connect::plan::JoinType;
 use spark_connect::udf::{eval_type, CommonInlineUserDefinedFunctionExpression, PythonUDFPayload};
 
@@ -526,6 +526,83 @@ fn no_to_plan(obj: &Bound<'_, PyAny>) -> PyErr {
     ))
 }
 
+/// The `TypeError` of protobuf for a name that is not a string, which the plan of PySpark's
+/// `_joinAsOf` raises for a name in `on` or an as-of column that is not a string.
+fn proto_string_type_error(value: &Bound<'_, PyAny>) -> PyErr {
+    let repr = value
+        .repr()
+        .map(|repr| repr.to_string())
+        .unwrap_or_default();
+    let typ = value
+        .get_type()
+        .repr()
+        .map(|typ| typ.to_string())
+        .unwrap_or_default();
+    pyo3::exceptions::PyTypeError::new_err(format!(
+        "{repr} has type {typ}, but expected one of: (<class 'bytes'>, <class 'str'>)"
+    ))
+}
+
+/// The join condition of `_joinAsOf` for its `on`, as PySpark's `AsOfJoin` plan sends it: a
+/// name or a list of names (`using_columns`), or a Column or a list of Columns, which are
+/// joined on their conjunction (`join_expr`).
+fn as_of_join_condition(
+    on: &Bound<'_, PyAny>,
+) -> PyResult<(Option<spark_connect::column::Column>, Vec<String>)> {
+    if let Ok(column) = on.extract::<PyRef<PyColumn>>() {
+        return Ok((Some(column.column.clone()), Vec::new()));
+    }
+    if on.is_instance_of::<pyo3::types::PyString>() {
+        return Ok((None, vec![on.extract::<String>()?]));
+    }
+    // As PySpark's `_collect_references`, which iterates `on` first, and then its `to_plan`.
+    if !on.is_instance_of::<PyList>() {
+        on.try_iter()?;
+        return Err(no_to_plan(on));
+    }
+    let list = on.cast::<PyList>()?;
+    // An empty list is no condition at all.
+    let Some(first) = list.iter().next() else {
+        return Ok((None, Vec::new()));
+    };
+    if first.is_instance_of::<pyo3::types::PyString>() {
+        let names = list
+            .iter()
+            .map(|item| {
+                item.extract::<String>()
+                    .map_err(|_| proto_string_type_error(&item))
+            })
+            .collect::<PyResult<Vec<_>>>()?;
+        return Ok((None, names));
+    }
+    let mut condition: Option<spark_connect::column::Column> = None;
+    for item in list.iter() {
+        let Ok(col) = item.extract::<PyRef<PyColumn>>() else {
+            return Err(no_to_plan(&item));
+        };
+        condition = Some(match condition {
+            None => col.column.clone(),
+            Some(left) => left.and(col.column.clone()),
+        });
+    }
+    Ok((condition, Vec::new()))
+}
+
+/// The as-of column of `_joinAsOf` of `df`: a Column as it is, or the column of `df` that
+/// has this name, as PySpark's `df._col(name)`.
+fn as_of_column(
+    df: &DataFrame,
+    column: &Bound<'_, PyAny>,
+) -> PyResult<spark_connect::column::Column> {
+    if let Ok(column) = column.extract::<PyRef<PyColumn>>() {
+        return Ok(column.column.clone());
+    }
+    let name = column
+        .extract::<String>()
+        .map_err(|_| proto_string_type_error(column))?;
+    Ok(df.col(&name))
+}
+
 /// Resolve a filter/where condition: a Column, or a SQL string parsed via `expr`.
 fn cond_column(condition: &Bound<'_, PyAny>) -> PyResult<spark_connect::column::Column> {
     if let Ok(c) = condition.extract::<PyRef<PyColumn>>() {
@@ -941,6 +1018,62 @@ impl PyDataFrame {
         Err(no_to_plan(&on))
     }
 
+    /// `DataFrame._joinAsOf(other, leftAsOfColumn, rightAsOfColumn, on=None, how=None, *,
+    /// tolerance=None, allowExactMatches=True, direction="backward")`, the as-of join of
+    /// pandas-on-Spark's `merge_asof`, as PySpark's. The as-of columns are Columns, or names of
+    /// their DataFrame. `on` is a name, a Column condition, or a list of either. `how` (inner by
+    /// default) and `direction` are sent as they are, and the server checks them.
+    #[pyo3(name = "_joinAsOf", signature = (other, leftAsOfColumn, rightAsOfColumn, on=None, how=None, *, tolerance=None, allowExactMatches=true, direction="backward"))]
+    #[allow(non_snake_case, clippy::too_many_arguments)]
+    fn join_as_of(
+        &self,
+        other: &Bound<'_, PyAny>,
+        leftAsOfColumn: &Bound<'_, PyAny>,
+        rightAsOfColumn: &Bound<'_, PyAny>,
+        on: Option<Bound<'_, PyAny>>,
+        how: Option<Bound<'_, PyAny>>,
+        tolerance: Option<Bound<'_, PyAny>>,
+        allowExactMatches: bool,
+        direction: &str,
+    ) -> PyResult<PyDataFrame> {
+        // As PySpark's `_check_same_session`, first.
+        let other = self.same_session(other)?;
+        let join_type = match how.filter(|how| !how.is_none()) {
+            None => "inner".to_string(),
+            Some(how) => how.extract::<String>().map_err(|_| {
+                pyo3::exceptions::PyAssertionError::new_err("how should be a string")
+            })?,
+        };
+        let tolerance = match tolerance.filter(|tolerance| !tolerance.is_none()) {
+            None => None,
+            Some(tolerance) => {
+                let column = tolerance.extract::<PyRef<PyColumn>>().map_err(|_| {
+                    pyo3::exceptions::PyAssertionError::new_err("tolerance should be Column")
+                })?;
+                Some(column.column.clone())
+            }
+        };
+        let (on, using_columns) = match on.filter(|on| !on.is_none()) {
+            None => (None, Vec::new()),
+            Some(on) => as_of_join_condition(&on)?,
+        };
+        let options = AsOfJoinOptions {
+            left_as_of: as_of_column(&self.dataframe, leftAsOfColumn)?,
+            right_as_of: as_of_column(&other.dataframe, rightAsOfColumn)?,
+            on,
+            using_columns,
+            join_type,
+            tolerance,
+            allow_exact_matches: allowExactMatches,
+            direction: direction.to_string(),
+        };
+        let df = self
+            .dataframe
+            .as_of_join(&other.dataframe, options)
+            .to_pyerr()?;
+        Ok(PyDataFrame::new(df))
+    }
+
     /// Cross join.
     fn crossJoin(&self, other: &Bound<'_, PyAny>) -> PyResult<PyDataFrame> {
         let other = self.same_session(other)?;

crates/spark-connect/src/dataframe.rs

@@ -5,7 +5,7 @@
 use std::collections::HashMap;
 use std::sync::{Arc, Mutex};
 
-use spark_connect_core::error::{Result, SparkError};
+use spark_connect_core::error::{Result, SparkError, SparkErrorKind};
 use spark_connect_core::runtime::{block_on, get_runtime};
 use spark_connect_proto as proto;
 
@@ -41,6 +41,47 @@ pub(crate) type ExecutionSlot = Arc<Mutex<Option<ExecutionInfo>>>;
 
 /// Iterator over rows from a DataFrame, yielded lazily as the server streams results.
 ///
+/// The options of an as-of join, as `DataFrame._joinAsOf` takes them.
+///
+/// `on` and `using_columns` are the join condition: a condition, or names that both sides
+/// share. PySpark sets at most one of them. The defaults are those of PySpark: an inner join,
+/// exact matches allowed, and the `"backward"` direction.
+#[derive(Debug, Clone)]
+pub struct AsOfJoinOptions {
+    /// The as-of column of this DataFrame.
+    pub left_as_of: Column,
+    /// The as-of column of the other DataFrame.
+    pub right_as_of: Column,
+    /// The join condition, sent as `join_expr`.
+    pub on: Option<Column>,
+    /// Column names that both sides share, sent as `using_columns`.
+    pub using_columns: Vec<String>,
+    /// The join type, such as `"inner"` or `"left"`, sent as it is.
+    pub join_type: String,
+    /// The as-of tolerance: the largest distance a match may have.
+    pub tolerance: Option<Column>,
+    /// Whether a row may match a row with the same as-of value.
+    pub allow_exact_matches: bool,
+    /// `"backward"`, `"forward"` or `"nearest"`, sent as it is.
+    pub direction: String,
+}
+
+impl AsOfJoinOptions {
+    /// Options joining on `left_as_of` and `right_as_of`, with the defaults of PySpark.
+    pub fn new(left_as_of: Column, right_as_of: Column) -> Self {
+        Self {
+            left_as_of,
+            right_as_of,
+            on: None,
+            using_columns: Vec::new(),
+            join_type: "inner".to_string(),
+            tolerance: None,
+            allow_exact_matches: true,
+            direction: "backward".to_string(),
+        }
+    }
+}
+
 /// Returned by `DataFrame::to_local_iterator()`, this iterator consumes the ExecutePlan
 /// response stream incrementally and yields Row objects without buffering the entire result.
 pub struct LocalRowIterator {
@@ -532,6 +573,35 @@ impl DataFrame {
         DataFrame::new(self.session.clone(), plan)
     }
 
+    /// As-of join with another DataFrame: each row of this DataFrame is matched with the
+    /// nearest row of `right` on the as-of columns, as `DataFrame._joinAsOf` does.
+    ///
+    /// Errors with `SESSION_NOT_SAME` when `right` is of another session. The join type and
+    /// the direction are sent as they are; the server checks them.
+    pub fn as_of_join(&self, right: &DataFrame, options: AsOfJoinOptions) -> Result<DataFrame> {
+        // As PySpark's `_check_same_session`.
+        if self.session.client().session_id() != right.session.client().session_id() {
+            return Err(SparkError::classed(
+                SparkErrorKind::SessionNotSame,
+                "SESSION_NOT_SAME",
+                &[],
+            ));
+        }
+        let plan = LogicalPlan::AsOfJoin {
+            left: Arc::clone(&self.plan),
+            right: Arc::clone(&right.plan),
+            left_as_of: options.left_as_of,
+            right_as_of: options.right_as_of,
+            on: options.on,
+            using_columns: options.using_columns,
+            join_type: options.join_type,
+            tolerance: options.tolerance,
+            allow_exact_matches: options.allow_exact_matches,
+            direction: options.direction,
+        };
+        Ok(DataFrame::new(self.session.clone(), plan))
+    }
+
     /// Cross join.
     pub fn cross_join(&self, right: &DataFrame) -> DataFrame {
         self.join(right, None, JoinType::Cross)
@@ -4312,6 +4382,82 @@ mod plan_construction_tests {
         assert_eq!(joined.plan.to_proto(), expected);
     }
 
+    /// As PySpark's `AsOfJoin` plan: the as-of columns, the names, the join type, the tolerance,
+    /// whether exact matches are allowed, and the direction are each sent.
+    #[test]
+    fn as_of_join_sends_every_field_of_the_as_of_join() {
+        use crate::functions::col;
+        use proto::relation::RelType;
+
+        let spark = session();
+        let left = spark.range(3).unwrap();
+        let right = spark.range(4).unwrap();
+        let options = AsOfJoinOptions {
+            using_columns: vec!["id".to_string()],
+            join_type: "left".to_string(),
+            tolerance: Some(crate::column::lit(1)),
+            allow_exact_matches: false,
+            direction: "forward".to_string(),
+            ..AsOfJoinOptions::new(col("id"), col("id"))
+        };
+        let joined = left.as_of_join(&right, options).unwrap();
+        let relation = joined.plan.to_proto();
+        let Some(RelType::AsOfJoin(as_of)) = &relation.rel_type else {
+            panic!("expected AsOfJoin: {relation:?}");
+        };
+        assert_eq!(as_of.left.as_deref(), Some(&left.plan.to_proto()));
+        assert_eq!(as_of.right.as_deref(), Some(&right.plan.to_proto()));
+        assert_eq!(as_of.left_as_of, Some(col("id").to_proto()));
+        assert_eq!(as_of.right_as_of, Some(col("id").to_proto()));
+        assert_eq!(as_of.join_expr, None);
+        assert_eq!(as_of.using_columns, vec!["id".to_string()]);
+        assert_eq!(as_of.join_type, "left");
+        assert_eq!(as_of.tolerance, Some(crate::column::lit(1).to_proto()));
+        assert!(!as_of.allow_exact_matches);
+        assert_eq!(as_of.direction, "forward");
+    }
+
+    /// The defaults of PySpark's `_joinAsOf` (inner, exact matches, backward), and a condition
+    /// sent as `join_expr` instead of names.
+    #[test]
+    fn as_of_join_defaults_and_sends_a_condition_as_join_expr() {
+        let spark = session();
+        let left = spark.range(3).unwrap();
+        let right = spark.range(4).unwrap();
+        let condition = left.col("id").eq(right.col("id"));
+        let options = AsOfJoinOptions {
+            on: Some(condition.clone()),
+            ..AsOfJoinOptions::new(left.col("id"), right.col("id"))
+        };
+        let joined = left.as_of_join(&right, options).unwrap();
+        let relation = joined.plan.to_proto();
+        let Some(proto::relation::RelType::AsOfJoin(as_of)) = &relation.rel_type else {
+            panic!("expected AsOfJoin: {relation:?}");
+        };
+        assert_eq!(as_of.join_expr, Some(condition.to_proto()));
+        assert!(as_of.using_columns.is_empty());
+        assert_eq!(as_of.join_type, "inner");
+        assert_eq!(as_of.tolerance, None);
+        assert!(as_of.allow_exact_matches);
+        assert_eq!(as_of.direction, "backward");
+    }
+
+    /// As PySpark's `_check_same_session`: a DataFrame of another session is not joined.
+    #[test]
+    fn as_of_join_rejects_a_dataframe_of_another_session() {
+        use crate::functions::col;
+
+        let spark = session();
+        let left = spark.range(3).unwrap();
+        let right = spark.new_session().range(3).unwrap();
+        let options = AsOfJoinOptions::new(col("id"), col("id"));
+        let error = left
+            .as_of_join(&right, options)
+            .err()
+            .expect("a session mismatch");
+        assert_eq!(error.get_error_class().as_deref(), Some("SESSION_NOT_SAME"));
+    }
+
     #[test]
     fn observed_metrics_set_the_observations_of_their_relations() {
         use crate::observation::Observation;

crates/spark-connect/src/plan.rs

@@ -442,6 +442,19 @@ pub enum LogicalPlan {
         other_grouping_expressions: Vec<Expression>,
         func: CommonInlineUserDefinedFunctionExpression,
     },
+    /// AsOfJoin: `DataFrame._joinAsOf`, which pandas-on-Spark's `merge_asof` calls.
+    AsOfJoin {
+        left: Arc<LogicalPlan>,
+        right: Arc<LogicalPlan>,
+        left_as_of: Column,
+        right_as_of: Column,
+        on: Option<Column>,
+        using_columns: Vec<String>,
+        join_type: String,
+        tolerance: Option<Column>,
+        allow_exact_matches: bool,
+        direction: String,
+    },
     /// NearestByJoin: `DataFrame.nearestByJoin`.
     NearestByJoin {
         left: Arc<LogicalPlan>,
@@ -821,6 +834,33 @@ impl LogicalPlan {
                 relation.rel_type = Some(proto::relation::RelType::Join(Box::new(join)));
             }
 
+            LogicalPlan::AsOfJoin {
+                left,
+                right,
+                left_as_of,
+                right_as_of,
+                on,
+                using_columns,
+                join_type,
+                tolerance,
+                allow_exact_matches,
+                direction,
+            } => {
+                let aj = proto::AsOfJoin {
+                    left: Some(Box::new(relation_of(left, plan_ids))),
+                    right: Some(Box::new(relation_of(right, plan_ids))),
+                    left_as_of: Some(left_as_of.to_proto()),
+                    right_as_of: Some(right_as_of.to_proto()),
+                    join_expr: on.as_ref().map(|condition| condition.to_proto()),
+                    using_columns: using_columns.clone(),
+                    join_type: join_type.clone(),
+                    tolerance: tolerance.as_ref().map(|tolerance| tolerance.to_proto()),
+                    allow_exact_matches: *allow_exact_matches,
+                    direction: direction.clone(),
+                };
+                relation.rel_type = Some(proto::relation::RelType::AsOfJoin(Box::new(aj)));
+            }
+
             LogicalPlan::LateralJoin {
                 left,
                 right,