more tests

pola-rs · Oct 3, 2024 · ab152e6 · ab152e6
1 parent 71b8274
commit ab152e6
Show file tree

Hide file tree

Showing 2 changed files with 107 additions and 39 deletions.
diff --git a/tests/testthat/test-dataframe.R b/tests/testthat/test-dataframe.R
@@ -1764,43 +1764,77 @@ test_that("$to_dummies() works", {
 
 test_that("inequality joins work", {
   east = pl$DataFrame(
-   id = c(100, 101, 102),
-   dur = c(120, 140, 160),
-   rev = c(12, 14, 16),
-   cores = c(2, 8, 4)
- )
- west = pl$DataFrame(
-   t_id = c(404, 498, 676, 742),
-   time = c(90, 130, 150, 170),
-   cost = c(9, 13, 15, 16),
-   cores = c(4, 2, 1, 4)
- )
- out = east$join_where(
-   west,
-   pl$col("dur") < pl$col("time"),
-   pl$col("rev") < pl$col("cost")
- )
-
- expect_identical(
-   out$to_data_frame(),
-   data.frame(
-     id = rep(c(100, 101), 3:2),
-     dur = rep(c(120, 140), 3:2),
-     rev = rep(c(12, 14), 3:2),
-     cores = rep(c(2, 8), 3:2),
-     t_id = c(498, 676, 742, 676, 742),
-     time = c(130, 150, 170, 150, 170),
-     cost = c(13, 15, 16, 15, 16),
-     cores_right = c(2, 1, 4, 1, 4)
-   )
- )
-
- expect_error(
-   east$join_where(
-     west$lazy(),
-     pl$col("dur") < pl$col("time"),
-     pl$col("rev") < pl$col("cost")
-   ),
-   "`other` must be a DataFrame"
- )
+    id = c(100, 101, 102),
+    dur = c(120, 140, 160),
+    rev = c(12, 14, 16),
+    cores = c(2, 8, 4)
+  )
+  west = pl$DataFrame(
+    t_id = c(404, 498, 676, 742),
+    time = c(90, 130, 150, 170),
+    cost = c(9, 13, 15, 16),
+    cores = c(4, 2, 1, 4)
+  )
+  out = east$join_where(
+    west,
+    pl$col("dur") < pl$col("time"),
+    pl$col("rev") < pl$col("cost")
+  )
+
+  expect_identical(
+    out$to_data_frame(),
+    data.frame(
+      id = rep(c(100, 101), 3:2),
+      dur = rep(c(120, 140), 3:2),
+      rev = rep(c(12, 14), 3:2),
+      cores = rep(c(2, 8), 3:2),
+      t_id = c(498, 676, 742, 676, 742),
+      time = c(130, 150, 170, 150, 170),
+      cost = c(13, 15, 16, 15, 16),
+      cores_right = c(2, 1, 4, 1, 4)
+    )
+  )
+
+  expect_error(
+    east$join_where(
+      west$lazy(),
+      pl$col("dur") < pl$col("time"),
+      pl$col("rev") < pl$col("cost")
+    ),
+    "`other` must be a DataFrame"
+  )
+})
+
+test_that("inequality joins require suffix when identical column names", {
+  east = pl$DataFrame(
+    id = c(100, 101, 102),
+    dur = c(120, 140, 160),
+    rev = c(12, 14, 16),
+    cores = c(2, 8, 4)
+  )
+  west = pl$DataFrame(
+    t_id = c(404, 498, 676, 742),
+    dur = c(90, 130, 150, 170),
+    rev = c(9, 13, 15, 16),
+    cores = c(4, 2, 1, 4)
+  )
+  out = east$join_where(
+    west,
+    pl$col("dur") < pl$col("dur_right"),
+    pl$col("rev") < pl$col("rev_right")
+  )
+
+  expect_identical(
+    out$to_data_frame(),
+    data.frame(
+      id = rep(c(100, 101), 3:2),
+      dur = rep(c(120, 140), 3:2),
+      rev = rep(c(12, 14), 3:2),
+      cores = rep(c(2, 8), 3:2),
+      t_id = c(498, 676, 742, 676, 742),
+      dur_right = c(130, 150, 170, 150, 170),
+      rev_right = c(13, 15, 16, 15, 16),
+      cores_right = c(2, 1, 4, 1, 4)
+    )
+  )
 })
diff --git a/tests/testthat/test-lazy.R b/tests/testthat/test-lazy.R
@@ -1241,3 +1241,37 @@ test_that("inequality joins work", {
     "`other` must be a LazyFrame"
   )
 })
+
+test_that("inequality joins require suffix when identical column names", {
+  east = pl$LazyFrame(
+    id = c(100, 101, 102),
+    dur = c(120, 140, 160),
+    rev = c(12, 14, 16),
+    cores = c(2, 8, 4)
+  )
+  west = pl$LazyFrame(
+    t_id = c(404, 498, 676, 742),
+    dur = c(90, 130, 150, 170),
+    rev = c(9, 13, 15, 16),
+    cores = c(4, 2, 1, 4)
+  )
+  out = east$join_where(
+    west,
+    pl$col("dur") < pl$col("dur_right"),
+    pl$col("rev") < pl$col("rev_right")
+  )$collect()
+
+  expect_identical(
+    out$to_data_frame(),
+    data.frame(
+      id = rep(c(100, 101), 3:2),
+      dur = rep(c(120, 140), 3:2),
+      rev = rep(c(12, 14), 3:2),
+      cores = rep(c(2, 8), 3:2),
+      t_id = c(498, 676, 742, 676, 742),
+      dur_right = c(130, 150, 170, 150, 170),
+      rev_right = c(13, 15, 16, 15, 16),
+      cores_right = c(2, 1, 4, 1, 4)
+    )
+  )
+})