feat: add checkSortedness option to joinAsOf (#374)

Context: pola-rs/polars#21724

Adds `checkSortedness` option so that warnings can be suppressed when `by` options are provided, since polars is unable to check sortedness.
Keeps the default behavior of `true`.

Author: b-twice

__tests__/dataframe.test.ts

@@ -1933,6 +1933,24 @@ describe("join", () => {
     expect(out.shape).toEqual({ height: 15, width: 4 });
   });
 });
+describe("joinAsOf", () => {
+  const df = pl.DataFrame({ a: [1, 1, 1, 2, 2, 2], b: [2, 1, 3, 1, 2, 3] });
+  test("errors when not sorted", () => {
+    expect(() => df.joinAsof(df, { on: "b" })).toThrow(/sorted/i);
+  });
+
+  test("does not error when not sorted but by is specified", () => {
+    expect(() => df.joinAsof(df, { on: "b", by: "a" })).not.toThrow();
+  });
+
+  test("skips sortedness check when checkSortedness is false", () => {
+    const df = pl.DataFrame({ a: [1, 1, 1, 2, 2, 2], b: [2, 1, 3, 1, 2, 3] });
+
+    expect(() =>
+      df.joinAsof(df, { on: "b", checkSortedness: false }),
+    ).not.toThrow();
+  });
+});
 describe("io", () => {
   const df = pl.DataFrame([
     pl.Series("foo", [1, 2, 9], pl.Int16),

polars/dataframe.ts

@@ -990,6 +990,10 @@ export interface DataFrame<S extends Schema = any>
    *   - "3d12h4m25s" # 3 days, 12 hours, 4 minutes, and 25 seconds
    * @param options.allowParallel Allow the physical plan to optionally evaluate the computation of both DataFrames up to the join in parallel.
    * @param options.forceParallel Force the physical plan to evaluate the computation of both DataFrames up to the join in parallel.
+   * @param options.checkSortedness
+   *    Check the sortedness of the asof keys. If the keys are not sorted Polars
+   *    will error, or in case of 'by' argument raise a warning. This might become
+   *    a hard error in the future.
    *
    * @example
    * ```
@@ -1045,6 +1049,7 @@ export interface DataFrame<S extends Schema = any>
       tolerance?: number | string;
       allowParallel?: boolean;
       forceParallel?: boolean;
+      checkSortedness?: boolean;
     },
   ): DataFrame;
   lazy(): LazyDataFrame<S>;

polars/lazy/dataframe.ts

@@ -362,6 +362,10 @@ export interface LazyDataFrame<S extends Schema = any>
         - "3d12h4m25s" # 3 days, 12 hours, 4 minutes, and 25 seconds
       @param options.allowParallel Allow the physical plan to optionally evaluate the computation of both DataFrames up to the join in parallel.
       @param options.forceParallel Force the physical plan to evaluate the computation of both DataFrames up to the join in parallel.
+      @param options.checkSortedness 
+        Check the sortedness of the asof keys. If the keys are not sorted Polars
+        will error, or in case of 'by' argument raise a warning. This might become
+        a hard error in the future.
 
 
       @example
@@ -418,6 +422,7 @@ export interface LazyDataFrame<S extends Schema = any>
       tolerance?: number | string;
       allowParallel?: boolean;
       forceParallel?: boolean;
+      checkSortedness?: boolean;
     },
   ): LazyDataFrame;
   /**
@@ -1046,9 +1051,16 @@ export const _LazyDataFrame = (_ldf: any): LazyDataFrame => {
         allowParallel: true,
         forceParallel: false,
         strategy: "backward",
+        checkSortedness: true,
         ...options,
       };
-      const { suffix, strategy, allowParallel, forceParallel } = options;
+      const {
+        suffix,
+        strategy,
+        allowParallel,
+        forceParallel,
+        checkSortedness,
+      } = options;
       let leftOn: string | undefined;
       let rightOn: string | undefined;
       if (!other?._ldf) {
@@ -1105,6 +1117,7 @@ export const _LazyDataFrame = (_ldf: any): LazyDataFrame => {
         strategy,
         toleranceNum,
         toleranceStr,
+        checkSortedness ?? true,
       );
 
       return _LazyDataFrame(ldf);

src/lazy/dataframe.rs

@@ -312,6 +312,7 @@ impl JsLazyFrame {
         strategy: String,
         tolerance: Option<Wrap<AnyValue<'_>>>,
         tolerance_str: Option<String>,
+        check_sortedness: bool,
     ) -> JsLazyFrame {
         let strategy = match strategy.as_ref() {
             "forward" => AsofStrategy::Forward,
@@ -340,7 +341,7 @@ impl JsLazyFrame {
                 }),
                 tolerance_str: tolerance_str.map(|s| s.into()),
                 allow_eq: true,
-                check_sortedness: true,
+                check_sortedness,
             })))
             .suffix(suffix)
             .finish()