diff --git a/crates/polars-core/src/frame/mod.rs b/crates/polars-core/src/frame/mod.rs index a96c81812ba14..f447f46fa72a3 100644 --- a/crates/polars-core/src/frame/mod.rs +++ b/crates/polars-core/src/frame/mod.rs @@ -369,7 +369,7 @@ impl DataFrame { /// | 3 | Patricia | /// +-----+----------+ /// ``` - pub fn with_row_count(&self, name: &str, offset: Option) -> PolarsResult { + pub fn with_row_number(&self, name: &str, offset: Option) -> PolarsResult { let mut columns = Vec::with_capacity(self.columns.len() + 1); let offset = offset.unwrap_or(0); diff --git a/crates/polars-lazy/src/frame/mod.rs b/crates/polars-lazy/src/frame/mod.rs index 000181aa8fa57..039162852adae 100644 --- a/crates/polars-lazy/src/frame/mod.rs +++ b/crates/polars-lazy/src/frame/mod.rs @@ -1636,7 +1636,7 @@ impl LazyFrame { /// # Warning /// This can have a negative effect on query performance. This may for instance block /// predicate pushdown optimization. - pub fn with_row_count(mut self, name: &str, offset: Option) -> LazyFrame { + pub fn with_row_number(mut self, name: &str, offset: Option) -> LazyFrame { let add_row_count_in_map = match &mut self.logical_plan { LogicalPlan::Scan { file_options: options, diff --git a/crates/polars-lazy/src/scan/anonymous_scan.rs b/crates/polars-lazy/src/scan/anonymous_scan.rs index 2a26305eb84bf..e561f77890c56 100644 --- a/crates/polars-lazy/src/scan/anonymous_scan.rs +++ b/crates/polars-lazy/src/scan/anonymous_scan.rs @@ -42,7 +42,7 @@ impl LazyFrame { .into(); if let Some(rc) = args.row_count { - lf = lf.with_row_count(&rc.name, Some(rc.offset)) + lf = lf.with_row_number(&rc.name, Some(rc.offset)) }; Ok(lf) diff --git a/crates/polars-lazy/src/scan/file_list_reader.rs b/crates/polars-lazy/src/scan/file_list_reader.rs index 18a1a62a29d30..6b1841f322455 100644 --- a/crates/polars-lazy/src/scan/file_list_reader.rs +++ b/crates/polars-lazy/src/scan/file_list_reader.rs @@ -61,7 +61,7 @@ pub trait LazyFileListReader: Clone { lf = lf.slice(0, n_rows as IdxSize) }; if let Some(rc) = self.row_count() { - lf = lf.with_row_count(&rc.name, Some(rc.offset)) + lf = lf.with_row_number(&rc.name, Some(rc.offset)) }; Ok(lf) diff --git a/crates/polars-lazy/src/scan/ipc.rs b/crates/polars-lazy/src/scan/ipc.rs index db75efada68a4..85eef30a90cea 100644 --- a/crates/polars-lazy/src/scan/ipc.rs +++ b/crates/polars-lazy/src/scan/ipc.rs @@ -65,7 +65,7 @@ impl LazyFileListReader for LazyIpcReader { // it is a bit hacky, but this row_count function updates the schema if let Some(row_count) = args.row_count { - lf = lf.with_row_count(&row_count.name, Some(row_count.offset)) + lf = lf.with_row_number(&row_count.name, Some(row_count.offset)) } Ok(lf) diff --git a/crates/polars-lazy/src/scan/parquet.rs b/crates/polars-lazy/src/scan/parquet.rs index 253d839281944..a0dfcd2ab1c36 100644 --- a/crates/polars-lazy/src/scan/parquet.rs +++ b/crates/polars-lazy/src/scan/parquet.rs @@ -90,7 +90,7 @@ impl LazyFileListReader for LazyParquetReader { // it is a bit hacky, but this row_count function updates the schema if let Some(row_count) = row_count { - lf = lf.with_row_count(&row_count.name, Some(row_count.offset)) + lf = lf.with_row_number(&row_count.name, Some(row_count.offset)) } lf.opt_state.file_caching = true; diff --git a/crates/polars-lazy/src/tests/io.rs b/crates/polars-lazy/src/tests/io.rs index 19577d7bda61f..d8b666215adb5 100644 --- a/crates/polars-lazy/src/tests/io.rs +++ b/crates/polars-lazy/src/tests/io.rs @@ -373,42 +373,42 @@ fn test_row_count_on_files() -> PolarsResult<()> { for offset in [0 as IdxSize, 10] { let lf = LazyCsvReader::new(FOODS_CSV) .with_row_count(Some(RowCount { - name: "rc".into(), + name: "rn".into(), offset, })) .finish()?; assert!(row_count_at_scan(lf.clone())); let df = lf.collect()?; - let rc = df.column("rc")?; + let rn = df.column("rn")?; assert_eq!( - rc.idx()?.into_no_null_iter().collect::>(), + rn.idx()?.into_no_null_iter().collect::>(), (offset..27 + offset).collect::>() ); let lf = LazyFrame::scan_parquet(FOODS_PARQUET, Default::default())? - .with_row_count("rc", Some(offset)); + .with_row_number("rn", Some(offset)); assert!(row_count_at_scan(lf.clone())); let df = lf.collect()?; - let rc = df.column("rc")?; + let rn = df.column("rn")?; assert_eq!( - rc.idx()?.into_no_null_iter().collect::>(), + rn.idx()?.into_no_null_iter().collect::>(), (offset..27 + offset).collect::>() ); let lf = - LazyFrame::scan_ipc(FOODS_IPC, Default::default())?.with_row_count("rc", Some(offset)); + LazyFrame::scan_ipc(FOODS_IPC, Default::default())?.with_row_number("rn", Some(offset)); assert!(row_count_at_scan(lf.clone())); let df = lf.clone().collect()?; - let rc = df.column("rc")?; + let rn = df.column("rn")?; assert_eq!( - rc.idx()?.into_no_null_iter().collect::>(), + rn.idx()?.into_no_null_iter().collect::>(), (offset..27 + offset).collect::>() ); let out = lf - .filter(col("rc").gt(lit(-1))) + .filter(col("rn").gt(lit(-1))) .select([col("calories")]) .collect()?; assert!(out.column("calories").is_ok()); diff --git a/crates/polars-lazy/src/tests/optimization_checks.rs b/crates/polars-lazy/src/tests/optimization_checks.rs index cdc3e7e49e7ee..96d5641c6beb2 100644 --- a/crates/polars-lazy/src/tests/optimization_checks.rs +++ b/crates/polars-lazy/src/tests/optimization_checks.rs @@ -351,11 +351,11 @@ fn test_with_row_count_opts() -> PolarsResult<()> { let out = df .clone() .lazy() - .with_row_count("row_nr", None) + .with_row_number("row_number", None) .tail(5) .collect()?; let expected = df![ - "row_nr" => [5 as IdxSize, 6, 7, 8, 9], + "row_number" => [5 as IdxSize, 6, 7, 8, 9], "a" => [5, 6, 7, 8, 9], ]?; @@ -363,11 +363,11 @@ fn test_with_row_count_opts() -> PolarsResult<()> { let out = df .clone() .lazy() - .with_row_count("row_nr", None) + .with_row_number("row_number", None) .slice(1, 2) .collect()?; assert_eq!( - out.column("row_nr")? + out.column("row_number")? .idx()? .into_no_null_iter() .collect::>(), @@ -377,11 +377,11 @@ fn test_with_row_count_opts() -> PolarsResult<()> { let out = df .clone() .lazy() - .with_row_count("row_nr", None) + .with_row_number("row_number", None) .filter(col("a").eq(lit(3i32))) .collect()?; assert_eq!( - out.column("row_nr")? + out.column("row_number")? .idx()? .into_no_null_iter() .collect::>(), @@ -392,10 +392,10 @@ fn test_with_row_count_opts() -> PolarsResult<()> { .clone() .lazy() .slice(1, 2) - .with_row_count("row_nr", None) + .with_row_number("row_number", None) .collect()?; assert_eq!( - out.column("row_nr")? + out.column("row_number")? .idx()? .into_no_null_iter() .collect::>(), @@ -405,10 +405,10 @@ fn test_with_row_count_opts() -> PolarsResult<()> { let out = df .lazy() .filter(col("a").eq(lit(3i32))) - .with_row_count("row_nr", None) + .with_row_number("row_number", None) .collect()?; assert_eq!( - out.column("row_nr")? + out.column("row_number")? .idx()? .into_no_null_iter() .collect::>(), diff --git a/crates/polars-lazy/src/tests/projection_queries.rs b/crates/polars-lazy/src/tests/projection_queries.rs index 0ef0646727139..da08a61d04332 100644 --- a/crates/polars-lazy/src/tests/projection_queries.rs +++ b/crates/polars-lazy/src/tests/projection_queries.rs @@ -65,7 +65,7 @@ fn test_cross_join_pd() -> PolarsResult<()> { } #[test] -fn test_row_count_pd() -> PolarsResult<()> { +fn test_row_number_pd() -> PolarsResult<()> { let df = df![ "x" => [1, 2, 3], "y" => [3, 2, 1], @@ -73,12 +73,12 @@ fn test_row_count_pd() -> PolarsResult<()> { let df = df .lazy() - .with_row_count("row_count", None) - .select([col("row_count"), col("x") * lit(3i32)]) + .with_row_number("row_number", None) + .select([col("row_number"), col("x") * lit(3i32)]) .collect()?; let expected = df![ - "row_count" => [0 as IdxSize, 1, 2], + "row_number" => [0 as IdxSize, 1, 2], "x" => [3i32, 6, 9] ]?; diff --git a/crates/polars-plan/src/logical_plan/functions/mod.rs b/crates/polars-plan/src/logical_plan/functions/mod.rs index 7517d2a1f627c..c3fbaf718e548 100644 --- a/crates/polars-plan/src/logical_plan/functions/mod.rs +++ b/crates/polars-plan/src/logical_plan/functions/mod.rs @@ -347,7 +347,7 @@ impl FunctionNode { let args = (**args).clone(); df.melt2(args) }, - RowCount { name, offset, .. } => df.with_row_count(name.as_ref(), *offset), + RowCount { name, offset, .. } => df.with_row_number(name.as_ref(), *offset), } } } diff --git a/crates/polars/tests/it/lazy/explodes.rs b/crates/polars/tests/it/lazy/explodes.rs index 01cc6ff69db75..27c987e9a6610 100644 --- a/crates/polars/tests/it/lazy/explodes.rs +++ b/crates/polars/tests/it/lazy/explodes.rs @@ -10,9 +10,9 @@ fn test_explode_row_numbers() -> PolarsResult<()> { ]? .lazy() .select([col("text").str().split(lit(" ")).alias("tokens")]) - .with_row_count("row_nr", None) + .with_row_number("row_number", None) .explode([col("tokens")]) - .select([col("row_nr"), col("tokens")]) + .select([col("row_number"), col("tokens")]) .collect()?; assert_eq!(df.shape(), (8, 2)); diff --git a/crates/polars/tests/it/lazy/queries.rs b/crates/polars/tests/it/lazy/queries.rs index 90a576720a143..cb429b2bead40 100644 --- a/crates/polars/tests/it/lazy/queries.rs +++ b/crates/polars/tests/it/lazy/queries.rs @@ -142,13 +142,13 @@ fn test_sorted_path() -> PolarsResult<()> { let out = df .lazy() - .with_row_count("row_nr", None) + .with_row_number("row_number", None) .explode(["a"]) - .group_by(["row_nr"]) + .group_by(["row_number"]) .agg([col("a").count().alias("count")]) .collect()?; - let s = out.column("row_nr")?; + let s = out.column("row_number")?; assert_eq!(s.is_sorted_flag(), IsSorted::Ascending); Ok(()) diff --git a/docs/src/rust/user-guide/expressions/column-selections.rs b/docs/src/rust/user-guide/expressions/column-selections.rs index d33ed96531f30..ba1a99fa101d0 100644 --- a/docs/src/rust/user-guide/expressions/column-selections.rs +++ b/docs/src/rust/user-guide/expressions/column-selections.rs @@ -16,7 +16,7 @@ fn main() -> Result<(), Box> { "logged_at" => date_range("logged_at", NaiveDate::from_ymd_opt(2022, 1, 1).unwrap().and_hms_opt(0, 0, 0).unwrap(), NaiveDate::from_ymd_opt(2022, 1, 1).unwrap().and_hms_opt(0, 0, 2).unwrap(), Duration::parse("1s"),ClosedWindow::Both, TimeUnit::Milliseconds, None)?, )? - .with_row_count("rn", None)?; + .with_row_number("rn", None)?; println!("{}", &df); // --8<-- [end:selectors_df] diff --git a/py-polars/src/dataframe.rs b/py-polars/src/dataframe.rs index 480eae5a9ed93..760f0360b1f1f 100644 --- a/py-polars/src/dataframe.rs +++ b/py-polars/src/dataframe.rs @@ -1139,7 +1139,7 @@ impl PyDataFrame { pub fn with_row_number(&self, name: &str, offset: Option) -> PyResult { let df = self .df - .with_row_count(name, offset) + .with_row_number(name, offset) .map_err(PyPolarsErr::from)?; Ok(df.into()) } diff --git a/py-polars/src/lazyframe/mod.rs b/py-polars/src/lazyframe/mod.rs index 181faa7e2b911..75172f4c91fbf 100644 --- a/py-polars/src/lazyframe/mod.rs +++ b/py-polars/src/lazyframe/mod.rs @@ -988,7 +988,7 @@ impl PyLazyFrame { fn with_row_number(&self, name: &str, offset: Option) -> Self { let ldf = self.ldf.clone(); - ldf.with_row_count(name, offset).into() + ldf.with_row_number(name, offset).into() } #[pyo3(signature = (lambda, predicate_pushdown, projection_pushdown, slice_pushdown, streamable, schema, validate_output))]