diff --git a/py-ltseq/ltseq/joins.py b/py-ltseq/ltseq/joins.py index b8427e9..782d471 100644 --- a/py-ltseq/ltseq/joins.py +++ b/py-ltseq/ltseq/joins.py @@ -480,12 +480,10 @@ def asof_join( try: joined_inner = left_table._inner.asof_join( right_table._inner, - left_time_col, - right_time_col, + (left_time_col, right_time_col), resolved, suffix, - left_by, - right_by, + list(zip(left_by, right_by)), ) except RuntimeError as e: raise RuntimeError(f"Asof join failed: {e}") diff --git a/py-ltseq/ltseq/ltseq_core.pyi b/py-ltseq/ltseq/ltseq_core.pyi index a04d3f4..8cd121d 100644 --- a/py-ltseq/ltseq/ltseq_core.pyi +++ b/py-ltseq/ltseq/ltseq_core.pyi @@ -55,7 +55,7 @@ class LTSeqTable: def join(self, other: LTSeqTable, left_key_expr: Any, right_key_expr: Any, how: str, suffix: str) -> LTSeqTable: ... def join_prefixed(self, other: LTSeqTable, left_key_expr: Any, right_key_expr: Any, how: str, alias: str) -> LTSeqTable: ... def lookup_join(self, other: LTSeqTable, key: str, target_key: str, alias: str) -> LTSeqTable: ... - def asof_join(self, other: LTSeqTable, left_key: str, right_key: str, direction: str, suffix: str, left_by_cols: list[str] = ..., right_by_cols: list[str] = ...) -> LTSeqTable: ... + def asof_join(self, other: LTSeqTable, on: tuple[str, str], direction: str, suffix: str, by: list[tuple[str, str]] = ...) -> LTSeqTable: ... def semi_join(self, other: LTSeqTable, left_key_expr: Any, right_key_expr: Any) -> LTSeqTable: ... def anti_join(self, other: LTSeqTable, left_key_expr: Any, right_key_expr: Any) -> LTSeqTable: ... def agg(self, grouping_expr: dict[str, Any] | None, agg_exprs: dict[str, dict[str, Any]]) -> LTSeqTable: ... diff --git a/src/lib.rs b/src/lib.rs index 819a8d7..6414afe 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -982,10 +982,10 @@ impl LTSeqTable { /// # Arguments /// /// * `other` - The right table to join with - /// * `left_time_col` - Column name for time/key in left table - /// * `right_time_col` - Column name for time/key in right table + /// * `on` - `(left, right)` time/key column names /// * `direction` - Match direction: "backward", "forward", or "nearest" - /// * `alias` - Prefix for right table columns (e.g., "_other") + /// * `suffix` - Appended to a right column whose name collides with a left one + /// * `by` - `(left, right)` pairs of columns that must also be equal /// /// # Direction Semantics /// @@ -1002,33 +1002,27 @@ impl LTSeqTable { /// /// Uses binary search for O(N log M) complexity where N = left rows, M = right rows. /// Both tables must be sorted by their respective time columns. - #[expect( - clippy::too_many_arguments, - reason = "arg count is dictated by the exposed Python signature below (see #150)" - )] - #[pyo3(signature = (other, left_time_col, right_time_col, direction, suffix, left_by_cols=Vec::new(), right_by_cols=Vec::new()))] + #[pyo3(signature = (other, on, direction, suffix, by=Vec::new()))] fn asof_join( &self, py: Python<'_>, other: <SeqTable, - left_time_col: &str, - right_time_col: &str, + on: (String, String), direction: &str, suffix: &str, - left_by_cols: Vec, - right_by_cols: Vec, + by: Vec<(String, String)>, ) -> PyResult { + let (left_by_cols, right_by_cols) = by.into_iter().unzip(); + let options = crate::ops::asof_join::AsofJoinOptions { + left_time_col: &on.0, + right_time_col: &on.1, + direction, + suffix, + left_by_cols, + right_by_cols, + }; detached(py, || { - crate::ops::asof_join::asof_join_impl( - self, - other, - left_time_col, - right_time_col, - direction, - suffix, - left_by_cols, - right_by_cols, - ) + crate::ops::asof_join::asof_join_impl(self, other, options) }) } diff --git a/src/ops/asof_join.rs b/src/ops/asof_join.rs index b9b59f9..a69a2dd 100644 --- a/src/ops/asof_join.rs +++ b/src/ops/asof_join.rs @@ -105,6 +105,25 @@ fn by_col_indices( .collect() } +/// How an as-of join matches: the arguments of the PyO3 `asof_join` binding +/// besides the two tables. The binding takes its column bindings as pairs, +/// `on: (left, right)` and `by: [(left, right), ...]`, and unpacks them here +/// into per-side fields, unvalidated; `asof_join_impl` validates them. +pub struct AsofJoinOptions<'a> { + /// Time/key column of the left table + pub left_time_col: &'a str, + /// Time/key column of the right table + pub right_time_col: &'a str, + /// `"backward"`, `"forward"` or `"nearest"` + pub direction: &'a str, + /// Appended to a right column whose name collides with a left one + pub suffix: &'a str, + /// Left `by=` columns, matched pairwise with `right_by_cols` + pub left_by_cols: Vec, + /// Right `by=` columns + pub right_by_cols: Vec, +} + /// As-of Join: Time-series join matching each left row with nearest right row /// /// For each row in the left table, finds the "nearest" matching row in the right table @@ -112,21 +131,20 @@ fn by_col_indices( /// /// Runs with the GIL released (`lib.rs::asof_join` wraps it in `gil::detached`): /// both collects and the binary-search matching take only plain Rust values. -#[expect( - clippy::too_many_arguments, - reason = "mirrors the flat parameter list of the PyO3 `asof_join` binding in lib.rs; \ - packing into a struct is a behavior-neutral refactor out of scope for #150" -)] pub fn asof_join_impl( table: <SeqTable, other: <SeqTable, - left_time_col: &str, - right_time_col: &str, - direction: &str, - suffix: &str, - left_by_cols: Vec, - right_by_cols: Vec, + options: AsofJoinOptions<'_>, ) -> Result { + let AsofJoinOptions { + left_time_col, + right_time_col, + direction, + suffix, + left_by_cols, + right_by_cols, + } = options; + // 1. Validate both tables have data let (_, stored_schema_left) = table.require_df_and_schema()?; let (_, stored_schema_right) = other.require_df_and_schema()?; diff --git a/src/ops/join.rs b/src/ops/join.rs index d5246b3..a1747ec 100644 --- a/src/ops/join.rs +++ b/src/ops/join.rs @@ -254,22 +254,26 @@ fn join_on_keys( .map_err(|e| LtseqError::Runtime(format!("Native join failed: {}", e))) } +/// The right table of a suffixed join, its colliding columns renamed. +struct RenamedRight { + /// The right DataFrame with the renamed columns. + df: DataFrame, + /// `(old, new)` name of every right column, in the right schema's order; + /// `old == new` for a column that did not collide. + rename_map: Vec<(String, String)>, + /// The right join key columns under their new names. + key_cols: Vec, +} + /// Rename right-table columns that collide with the left table by appending -/// `suffix` (Polars semantics: only conflicting columns are renamed). Returns -/// the renamed DataFrame, the `(old, new)` rename map, and the renamed names of -/// the join key columns. -#[expect( - clippy::type_complexity, - reason = "returns (renamed df, old→new rename pairs, right key cols after rename); \ - a named struct for one internal call site is out of scope for #150" -)] +/// `suffix` (Polars semantics: only conflicting columns are renamed). fn rename_right_df_for_join( df: DataFrame, left_schema: &ArrowSchema, right_schema: &ArrowSchema, suffix: &str, right_key_cols: &[String], -) -> PyResult<(DataFrame, Vec<(String, String)>, Vec)> { +) -> PyResult { let rename_map = right_rename_map(left_schema, right_schema, suffix)?; let select_exprs: Vec = rename_map @@ -298,7 +302,11 @@ fn rename_right_df_for_join( .select(select_exprs) .map_err(|e| LtseqError::Runtime(format!("Failed to rename right columns: {}", e)))?; - Ok((renamed_df, rename_map, new_key_cols)) + Ok(RenamedRight { + df: renamed_df, + rename_map, + key_cols: new_key_cols, + }) } /// Build result LTSeqTable from a DataFrame. @@ -360,7 +368,11 @@ pub fn join_impl( }; // Rename only the right columns that collide with the left table (suffix). - let (renamed_right_df, rename_map, new_right_keys) = rename_right_df_for_join( + let RenamedRight { + df: renamed_right_df, + rename_map, + key_cols: new_right_keys, + } = rename_right_df_for_join( (**df_right).clone(), stored_schema_left, stored_schema_right, diff --git a/src/ops/parallel_scan.rs b/src/ops/parallel_scan.rs index 6b288af..ae2e67b 100644 --- a/src/ops/parallel_scan.rs +++ b/src/ops/parallel_scan.rs @@ -307,6 +307,21 @@ struct RgBoundaryInfo { one_partition: bool, } +/// What every row group's fused read+match shares: the file, the columns +/// read from it, and the pattern to match. Built once per file and lent to +/// each rayon task, which adds only its row group's index. +struct RowGroupScan<'a> { + parquet_path: &'a str, + projection_mask: &'a ProjectionMask, + /// Position of the partition column in the projected schema + partition_col_idx: usize, + step_predicates: &'a [PyExpr], + fast_path_plan: Option<&'a StartsWithFastPathPlan>, + num_steps: usize, + /// Column name to position in the projected schema + name_to_idx: &'a HashMap, +} + /// Result from fused read+match of a single row group. struct RgMatchResult { /// Number of intra-RG pattern matches found @@ -408,22 +423,19 @@ pub fn parallel_pattern_match_count( // Step 3+4: Fused read + match per RG in parallel. // Each task reads one RG, pattern-matches it, extracts boundary data, drops RG data. // This eliminates the 1.4s dealloc overhead from holding all RG data in memory. - let path_str = parquet_path.to_string(); + let scan = RowGroupScan { + parquet_path, + projection_mask: &projection_mask, + partition_col_idx: part_col_idx, + step_predicates, + fast_path_plan: fast_path_plan.as_ref(), + num_steps, + name_to_idx: &name_to_idx, + }; let rg_results: Vec = (0..num_row_groups) .into_par_iter() - .map(|rg_idx| { - read_match_and_extract_boundary( - &path_str, - rg_idx, - &projection_mask, - part_col_idx, - step_predicates, - fast_path_plan.as_ref(), - num_steps, - &name_to_idx, - ) - }) + .map(|rg_idx| read_match_and_extract_boundary(&scan, rg_idx)) .collect::, _>>() .map_err(|e| LtseqError::Runtime(format!("Parallel fused read+match failed: {}", e)))?; @@ -455,20 +467,19 @@ pub fn parallel_pattern_match_count( /// RG data. This means each RG's ~20MB of data is freed immediately after /// processing, eliminating the 1.4s dealloc overhead from holding all 814 RGs /// (~2.7GB) in memory simultaneously. -#[expect( - clippy::too_many_arguments, - reason = "packing the args into a struct is a behavior-neutral refactor out of scope for #150" -)] fn read_match_and_extract_boundary( - parquet_path: &str, + scan: &RowGroupScan<'_>, row_group_idx: usize, - projection_mask: &ProjectionMask, - partition_col_idx: usize, - step_predicates: &[PyExpr], - fast_path_plan: Option<&StartsWithFastPathPlan>, - num_steps: usize, - name_to_idx: &HashMap, ) -> Result { + let RowGroupScan { + parquet_path, + projection_mask, + partition_col_idx, + step_predicates, + fast_path_plan, + num_steps, + name_to_idx, + } = *scan; let file = File::open(parquet_path) .map_err(|e| format!("Failed to open Parquet for RG {}: {}", row_group_idx, e))?;