feat: type cross join

pola-rs · Jan 26, 2025 · a0514ee · a0514ee
1 parent 49c0a0e
commit a0514ee
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 18 deletions.
diff --git a/__tests__/type.test.ts b/__tests__/type.test.ts
@@ -55,11 +55,15 @@ describe("type tests", () => {
         fl: Float64;
       }>
     >(a.join(b, { leftOn: "id", rightOn: ["age"] }));
-
-    const c = a.join(b, { on: ["id"], suffix: "Right" });
-
-    const d = DataFrame({ a: [1n, 2n], id: [3, 4] })
-      .join(a, { on: "id" })
-      .getColumn("id");
+    expectType<
+      DataFrame<{
+        id: Int64;
+        id_right: Int64;
+        age: Int64;
+        age_right: Int64;
+        name: PlString;
+        fl: Float64;
+      }>
+    >(a.join(b, { how: "cross" }));
   });
 });
diff --git a/polars/dataframe.ts b/polars/dataframe.ts
@@ -206,24 +206,25 @@ type JoinSchemas<
     [K_SUFFIXED in keyof S1 &
       Exclude<
         keyof S2,
-        Opt extends Pick<JoinOptions, "on">
-          ? ExtractJoinKeys<Opt["on"]>
-          : Opt extends Pick<JoinOptions, "leftOn" | "rightOn">
-            ? ExtractJoinKeys<Opt["rightOn"]>
-            : never
+        Opt extends { how: "cross" }
+          ? never
+          : Opt extends Pick<JoinOptions, "on">
+            ? ExtractJoinKeys<Opt["on"]>
+            : Opt extends Pick<JoinOptions, "leftOn" | "rightOn">
+              ? ExtractJoinKeys<Opt["rightOn"]>
+              : never
       > as `${K_SUFFIXED extends string ? K_SUFFIXED : never}${ExtractSuffix<Opt>}`]: K_SUFFIXED extends string
       ? S2[K_SUFFIXED]
       : never;
   }
 >;
 
-type JoinDataFrames<T1, T2, Opt extends JoinOptions> = T1 extends DataFrame<
-  infer T1Schema
->
-  ? T2 extends DataFrame<infer T2Schema>
-    ? DataFrame<JoinSchemas<T1Schema, T2Schema, Opt>>
-    : never
-  : never;
+type JoinDataFrames<T1, T2, Opt extends JoinOptions> =
+  T1 extends DataFrame<infer T1Schema>
+    ? T2 extends DataFrame<infer T2Schema>
+      ? DataFrame<JoinSchemas<T1Schema, T2Schema, Opt>>
+      : never
+    : never;
 
 /**
  * A DataFrame is a two-dimensional data structure that represents data as a table