don't allow predicate pushdown if compared column is being coerced (#…

…3457)

.github/workflows/create-py-mac-universal2-release.yaml

@@ -19,7 +19,7 @@ jobs:
       - name: Install latest Rust nightly
         uses: actions-rs/toolchain@v1
         with:
-          toolchain: nightly-2022-04-01
+          toolchain: nightly-2022-05-20
           override: true
           components: rustfmt, clippy
       - name: Setup universal2 targets for Rust

.github/workflows/create-py-release-windows-macos.yaml

@@ -37,7 +37,7 @@ jobs:
             rm py-polars/README.md
             cp README.md py-polars/README.md
             cd py-polars
-            rustup override set nightly-2022-04-01
+            rustup override set nightly-2022-05-20
             export RUSTFLAGS='-C target-feature=+fxsr,+sse,+sse2,+sse3'
             maturin publish \
             --no-sdist \

.github/workflows/test-python.yaml

@@ -30,7 +30,7 @@ jobs:
           pip install -r py-polars/build.requirements.txt
       - name: Run formatting checks
         run: |
-          cd py-polars && black --check . && blackdoc --check . && isort --check . && rustup override set nightly-2022-04-01 && cargo fmt --all -- --check && cd ..
+          cd py-polars && black --check . && blackdoc --check . && isort --check . && rustup override set nightly-2022-05-20 && cargo fmt --all -- --check && cd ..
       - name: Run linting
         run: |
           cd py-polars && flake8 && cd ..
@@ -40,7 +40,7 @@ jobs:
       - name: Run tests
         run: |
           export RUSTFLAGS="-C debuginfo=0"
-          cd py-polars && rustup override set nightly-2022-04-01 && make venv && make test-with-cov
+          cd py-polars && rustup override set nightly-2022-05-20 && make venv && make test-with-cov
           cargo clippy
       - name: Check doc examples
         run: |

.github/workflows/test-windows-python.yaml

@@ -28,7 +28,7 @@ jobs:
         shell: bash
         run: |
           export RUSTFLAGS="-C debuginfo=0"
-          cd py-polars && rustup override set nightly-2022-04-01 && make build-and-test-no-venv
+          cd py-polars && rustup override set nightly-2022-05-01 && make build-and-test-no-venv
           cargo clippy
       # test if we can import polars without any requirements
       - name: Import polars

Makefile

@@ -10,7 +10,7 @@ fmt_toml:
 
 coverage:
 	@bash -c "\
-		rustup override set nightly-2022-04-01; \
+		rustup override set nightly-2022-05-20; \
 		source <(cargo llvm-cov show-env --export-prefix); \
 		export CARGO_TARGET_DIR=\$$CARGO_LLVM_COV_TARGET_DIR; \
 		export CARGO_INCREMENTAL=1; \

polars/polars-core/src/lib.rs

@@ -114,6 +114,14 @@ pub(crate) static STRING_CACHE: Lazy<StringCache> = Lazy::new(Default::default);
 // utility for the tests to ensure a single thread can execute
 pub static SINGLE_LOCK: Lazy<Mutex<()>> = Lazy::new(|| Mutex::new(()));
 
+#[cfg(feature = "dtype-categorical")]
+pub fn with_string_cache<F: FnOnce() -> T, T>(func: F) -> T {
+    toggle_string_cache(true);
+    let out = func();
+    toggle_string_cache(false);
+    out
+}
+
 /// Use a global string cache for the Categorical Types.
 ///
 /// This is used to cache the string categories locally.

polars/polars-lazy/src/logical_plan/format.rs

@@ -231,8 +231,16 @@ impl fmt::Debug for Expr {
                 }
             }
             Cast {
-                expr, data_type, ..
-            } => write!(f, "{:?}.cast({:?})", expr, data_type),
+                expr,
+                data_type,
+                strict,
+            } => {
+                if *strict {
+                    write!(f, "{:?}.strict_cast({:?})", expr, data_type)
+                } else {
+                    write!(f, "{:?}.cast({:?})", expr, data_type)
+                }
+            }
             Ternary {
                 predicate,
                 truthy,

polars/polars-lazy/src/logical_plan/optimizer/predicate_pushdown/utils.rs

@@ -149,9 +149,9 @@ pub(super) fn project_other_column_is_predicate_pushdown_boundary(
             | AExpr::AnonymousFunction {options: FunctionOptions { collect_groups: ApplyOptions::ApplyGroups, .. }, ..}
             | AExpr::Function {options: FunctionOptions { collect_groups: ApplyOptions::ApplyGroups, .. }, ..}
             | AExpr::BinaryExpr {..}
-            | AExpr::Cast {data_type: DataType::Float32 | DataType::Float64, ..}
-            // cast may create nulls
-            | AExpr::Cast {strict: false, ..}
+            // casts may produce null values, change values etc.
+            // they can fail in myriad ways
+            | AExpr::Cast {..}
             // still need to investigate this one
             | AExpr::Explode {..}
             // A groupby needs all rows for aggregation

polars/tests/it/lazy/expressions/is_in.rs

@@ -1,7 +1,6 @@
 use super::*;
 
 #[test]
-#[cfg(feature = "is_in")]
 fn test_is_in() -> Result<()> {
     let df = df![
         "x" => [1, 2, 3],

polars/tests/it/lazy/expressions/mod.rs

@@ -1,6 +1,7 @@
 mod apply;
 mod arity;
 mod expand;
+#[cfg(feature = "is_in")]
 mod is_in;
 mod slice;
 mod window;

polars/tests/it/lazy/predicate_queries.rs

@@ -1,4 +1,5 @@
 use super::*;
+use polars_core::{with_string_cache, SINGLE_LOCK};
 
 #[test]
 fn test_predicate_after_renaming() -> Result<()> {
@@ -137,3 +138,29 @@ fn test_filter_block_join() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+#[cfg(all(feature = "is_in", feature = "dtype-categorical"))]
+fn test_is_in_categorical_3420() -> Result<()> {
+    let df = df![
+        "a" => ["a", "b", "c", "d", "e"],
+        "b" => [1, 2, 3, 4, 5]
+    ]?;
+
+    let _guard = SINGLE_LOCK.lock();
+
+    let out: Result<_> = with_string_cache(|| {
+        let s = Series::new("x", ["a", "b", "c"]).strict_cast(&DataType::Categorical(None))?;
+        df.lazy()
+            .with_column(col("a").strict_cast(DataType::Categorical(None)))
+            .filter(col("a").is_in(lit(s).alias("x")))
+            .collect()
+    });
+    let mut expected = df![
+        "a" => ["a", "b", "c"],
+        "b" => [1, 2, 3]
+    ]?;
+    expected.try_apply("a", |s| s.cast(&DataType::Categorical(None)))?;
+    assert!(out?.frame_equal(&expected));
+    Ok(())
+}