fix(rust, python): panic when max_len -1 is reached (#6494)

pola-rs · Jan 27, 2023 · 95de367 · 95de367
1 parent 861c26f
commit 95de367
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 2 deletions.
diff --git a/polars/polars-core/Cargo.toml b/polars/polars-core/Cargo.toml
@@ -98,6 +98,7 @@ parquet = ["arrow/io_parquet"]
 
 # scale to terabytes?
 bigidx = ["polars-arrow/bigidx"]
+python = []
 
 serde-lazy = ["serde", "polars-arrow/serde", "indexmap/serde"]
 

diff --git a/polars/polars-core/src/chunked_array/ops/chunkops.rs b/polars/polars-core/src/chunked_array/ops/chunkops.rs
@@ -72,7 +72,17 @@ impl<T: PolarsDataType> ChunkedArray<T> {
                 _ => chunks.iter().fold(0, |acc, arr| acc + arr.len()),
             }
         }
-        self.length = inner(&self.chunks) as IdxSize
+        self.length = inner(&self.chunks) as IdxSize;
+        #[cfg(feature = "python")]
+        assert!(
+            self.length < IdxSize::MAX,
+            "Polars' maximum length reached. Consider installing 'polars-u64-idx'."
+        );
+        #[cfg(not(feature = "python"))]
+        assert!(
+            self.length < IdxSize::MAX,
+            "Polars' maximum length reached. Consider compiling with 'bigidx' feature."
+        );
     }
 
     pub fn rechunk(&self) -> Self {

diff --git a/py-polars/Cargo.toml b/py-polars/Cargo.toml
@@ -21,7 +21,7 @@ libc = "0.2"
 ndarray = "0.15"
 numpy = "0.16"
 once_cell = "1"
-polars-core = { path = "../polars/polars-core", default-features = false }
+polars-core = { path = "../polars/polars-core", features = ["python"], default-features = false }
 polars-lazy = { path = "../polars/polars-lazy", features = ["python"], default-features = false }
 pyo3 = { version = "0.16", features = ["abi3-py37", "extension-module", "multiple-pymethods"] }
 pyo3-built = { version = "0.4", optional = true }