depr(python,rust!): Rename row_count_name/row_count_offset parame…

…ters in IO functions to `row_index_*` (pola-rs#13563)
MarcNuebel · Jan 9, 2024 · 130c48f · 130c48f
1 parent 8730ced
commit 130c48f
Show file tree

Hide file tree

Showing 57 changed files with 485 additions and 455 deletions.
diff --git a/crates/polars-core/src/frame/mod.rs b/crates/polars-core/src/frame/mod.rs
@@ -384,8 +384,8 @@ impl DataFrame {
         DataFrame::new(columns)
     }
 
-    /// Add a row count in place.
-    pub fn with_row_count_mut(&mut self, name: &str, offset: Option<IdxSize>) -> &mut Self {
+    /// Add a row index column in place.
+    pub fn with_row_index_mut(&mut self, name: &str, offset: Option<IdxSize>) -> &mut Self {
         let offset = offset.unwrap_or(0);
         let mut ca = IdxCa::from_vec(
             name,

diff --git a/crates/polars-io/src/csv/mod.rs b/crates/polars-io/src/csv/mod.rs
@@ -69,4 +69,4 @@ use crate::csv::read_impl::CoreReader;
 use crate::mmap::MmapBytesReader;
 use crate::predicates::PhysicalIoExpr;
 use crate::utils::{get_reader_bytes, resolve_homedir};
-use crate::{RowCount, SerReader, SerWriter};
+use crate::{RowIndex, SerReader, SerWriter};
diff --git a/crates/polars-io/src/csv/read.rs b/crates/polars-io/src/csv/read.rs
@@ -151,7 +151,7 @@ where
     quote_char: Option<u8>,
     skip_rows_after_header: usize,
     try_parse_dates: bool,
-    row_count: Option<RowCount>,
+    row_index: Option<RowIndex>,
     /// Aggregates chunk afterwards to a single chunk.
     rechunk: bool,
     raise_if_empty: bool,
@@ -173,9 +173,9 @@ where
         self
     }
 
-    /// Add a `row_count` column.
-    pub fn with_row_count(mut self, rc: Option<RowCount>) -> Self {
-        self.row_count = rc;
+    /// Add a row index column.
+    pub fn with_row_index(mut self, row_index: Option<RowIndex>) -> Self {
+        self.row_index = row_index;
         self
     }
 
@@ -417,7 +417,7 @@ impl<'a, R: MmapBytesReader + 'a> CsvReader<'a, R> {
             std::mem::take(&mut self.predicate),
             to_cast,
             self.skip_rows_after_header,
-            std::mem::take(&mut self.row_count),
+            std::mem::take(&mut self.row_index),
             self.try_parse_dates,
             self.raise_if_empty,
             self.truncate_ragged_lines,
@@ -603,7 +603,7 @@ where
             quote_char: Some(b'"'),
             skip_rows_after_header: 0,
             try_parse_dates: false,
-            row_count: None,
+            row_index: None,
             raise_if_empty: true,
             truncate_ragged_lines: false,
         }

diff --git a/crates/polars-io/src/csv/read_impl/batched_mmap.rs b/crates/polars-io/src/csv/read_impl/batched_mmap.rs
@@ -153,7 +153,7 @@ impl<'a> CoreReader<'a> {
             str_columns,
             projection,
             starting_point_offset,
-            row_count: self.row_count,
+            row_index: self.row_index,
             comment_prefix: self.comment_prefix,
             quote_char: self.quote_char,
             eol_char: self.eol_char,
@@ -181,7 +181,7 @@ pub struct BatchedCsvReaderMmap<'a> {
     str_columns: StringColumns,
     projection: Vec<usize>,
     starting_point_offset: Option<usize>,
-    row_count: Option<RowCount>,
+    row_index: Option<RowIndex>,
     comment_prefix: Option<CommentPrefix>,
     quote_char: Option<u8>,
     eol_char: u8,
@@ -255,16 +255,16 @@ impl<'a> BatchedCsvReaderMmap<'a> {
                     cast_columns(&mut df, &self.to_cast, false, self.ignore_errors)?;
 
                     update_string_stats(&self.str_capacities, &self.str_columns, &df)?;
-                    if let Some(rc) = &self.row_count {
-                        df.with_row_count_mut(&rc.name, Some(rc.offset));
+                    if let Some(rc) = &self.row_index {
+                        df.with_row_index_mut(&rc.name, Some(rc.offset));
                     }
                     Ok(df)
                 })
                 .collect::<PolarsResult<Vec<_>>>()
         })?;
         self.file_chunks.clear();
 
-        if self.row_count.is_some() {
+        if self.row_index.is_some() {
             update_row_counts2(&mut chunks, self.rows_read)
         }
         for df in &chunks {

diff --git a/crates/polars-io/src/csv/read_impl/batched_read.rs b/crates/polars-io/src/csv/read_impl/batched_read.rs
@@ -236,7 +236,7 @@ impl<'a> CoreReader<'a> {
             str_columns,
             projection,
             starting_point_offset,
-            row_count: self.row_count,
+            row_index: self.row_index,
             comment_prefix: self.comment_prefix,
             quote_char: self.quote_char,
             eol_char: self.eol_char,
@@ -264,7 +264,7 @@ pub struct BatchedCsvReaderRead<'a> {
     str_columns: StringColumns,
     projection: Vec<usize>,
     starting_point_offset: Option<usize>,
-    row_count: Option<RowCount>,
+    row_index: Option<RowIndex>,
     comment_prefix: Option<CommentPrefix>,
     quote_char: Option<u8>,
     eol_char: u8,
@@ -352,16 +352,16 @@ impl<'a> BatchedCsvReaderRead<'a> {
                     cast_columns(&mut df, &self.to_cast, false, self.ignore_errors)?;
 
                     update_string_stats(&self.str_capacities, &self.str_columns, &df)?;
-                    if let Some(rc) = &self.row_count {
-                        df.with_row_count_mut(&rc.name, Some(rc.offset));
+                    if let Some(rc) = &self.row_index {
+                        df.with_row_index_mut(&rc.name, Some(rc.offset));
                     }
                     Ok(df)
                 })
                 .collect::<PolarsResult<Vec<_>>>()
         })?;
         self.file_chunks.clear();
 
-        if self.row_count.is_some() {
+        if self.row_index.is_some() {
             update_row_counts2(&mut chunks, self.rows_read)
         }
         for df in &chunks {

diff --git a/crates/polars-io/src/csv/read_impl/mod.rs b/crates/polars-io/src/csv/read_impl/mod.rs
@@ -26,7 +26,7 @@ use crate::csv::{CsvEncoding, NullValues};
 use crate::mmap::ReaderBytes;
 use crate::predicates::PhysicalIoExpr;
 use crate::utils::update_row_counts;
-use crate::RowCount;
+use crate::RowIndex;
 
 pub(crate) fn cast_columns(
     df: &mut DataFrame,
@@ -116,7 +116,7 @@ pub(crate) struct CoreReader<'a> {
     missing_is_null: bool,
     predicate: Option<Arc<dyn PhysicalIoExpr>>,
     to_cast: Vec<Field>,
-    row_count: Option<RowCount>,
+    row_index: Option<RowIndex>,
     truncate_ragged_lines: bool,
 }
 
@@ -206,7 +206,7 @@ impl<'a> CoreReader<'a> {
         predicate: Option<Arc<dyn PhysicalIoExpr>>,
         to_cast: Vec<Field>,
         skip_rows_after_header: usize,
-        row_count: Option<RowCount>,
+        row_index: Option<RowIndex>,
         try_parse_dates: bool,
         raise_if_empty: bool,
         truncate_ragged_lines: bool,
@@ -306,7 +306,7 @@ impl<'a> CoreReader<'a> {
             missing_is_null,
             predicate,
             to_cast,
-            row_count,
+            row_index,
             truncate_ragged_lines,
         })
     }
@@ -552,8 +552,8 @@ impl<'a> CoreReader<'a> {
         // An empty file with a schema should return an empty DataFrame with that schema
         if bytes.is_empty() {
             let mut df = DataFrame::from(self.schema.as_ref());
-            if let Some(ref row_count) = self.row_count {
-                df.insert_column(0, Series::new_empty(&row_count.name, &IDX_DTYPE))?;
+            if let Some(ref row_index) = self.row_index {
+                df.insert_column(0, Series::new_empty(&row_index.name, &IDX_DTYPE))?;
             }
             return Ok(df);
         }
@@ -618,8 +618,8 @@ impl<'a> CoreReader<'a> {
                                     .collect::<PolarsResult<_>>()?,
                             );
                             let current_row_count = local_df.height() as IdxSize;
-                            if let Some(rc) = &self.row_count {
-                                local_df.with_row_count_mut(&rc.name, Some(rc.offset));
+                            if let Some(rc) = &self.row_index {
+                                local_df.with_row_index_mut(&rc.name, Some(rc.offset));
                             };
 
                             cast_columns(&mut local_df, &self.to_cast, false, self.ignore_errors)?;
@@ -638,7 +638,7 @@ impl<'a> CoreReader<'a> {
                     .collect::<PolarsResult<Vec<_>>>()
             })?;
             let mut dfs = flatten(&dfs, None);
-            if self.row_count.is_some() {
+            if self.row_index.is_some() {
                 update_row_counts(&mut dfs, 0)
             }
             accumulate_dataframes_vertical(dfs.into_iter().map(|t| t.0))
@@ -687,8 +687,8 @@ impl<'a> CoreReader<'a> {
                         }
 
                         cast_columns(&mut df, &self.to_cast, false, self.ignore_errors)?;
-                        if let Some(rc) = &self.row_count {
-                            df.with_row_count_mut(&rc.name, Some(rc.offset));
+                        if let Some(rc) = &self.row_index {
+                            df.with_row_index_mut(&rc.name, Some(rc.offset));
                         }
                         let n_read = df.height() as IdxSize;
                         Ok((df, n_read))
@@ -738,15 +738,15 @@ impl<'a> CoreReader<'a> {
                         };
 
                         cast_columns(&mut df, &self.to_cast, false, self.ignore_errors)?;
-                        if let Some(rc) = &self.row_count {
-                            df.with_row_count_mut(&rc.name, Some(rc.offset));
+                        if let Some(rc) = &self.row_index {
+                            df.with_row_index_mut(&rc.name, Some(rc.offset));
                         }
                         let n_read = df.height() as IdxSize;
                         (df, n_read)
                     });
                 }
             }
-            if self.row_count.is_some() {
+            if self.row_index.is_some() {
                 update_row_counts(&mut dfs, 0)
             }
             accumulate_dataframes_vertical(dfs.into_iter().map(|t| t.0))

diff --git a/crates/polars-io/src/ipc/ipc_file.rs b/crates/polars-io/src/ipc/ipc_file.rs
@@ -44,7 +44,7 @@ use super::{finish_reader, ArrowReader};
 use crate::mmap::MmapBytesReader;
 use crate::predicates::PhysicalIoExpr;
 use crate::prelude::*;
-use crate::RowCount;
+use crate::RowIndex;
 
 /// Read Arrows IPC format into a DataFrame
 ///
@@ -71,7 +71,7 @@ pub struct IpcReader<R: MmapBytesReader> {
     pub(super) n_rows: Option<usize>,
     pub(super) projection: Option<Vec<usize>>,
     pub(crate) columns: Option<Vec<String>>,
-    pub(super) row_count: Option<RowCount>,
+    pub(super) row_index: Option<RowIndex>,
     memmap: bool,
     metadata: Option<read::FileMetadata>,
     schema: Option<ArrowSchemaRef>,
@@ -127,9 +127,9 @@ impl<R: MmapBytesReader> IpcReader<R> {
         self
     }
 
-    /// Add a `row_count` column.
-    pub fn with_row_count(mut self, row_count: Option<RowCount>) -> Self {
-        self.row_count = row_count;
+    /// Add a row index column.
+    pub fn with_row_index(mut self, row_index: Option<RowIndex>) -> Self {
+        self.row_index = row_index;
         self
     }
 
@@ -173,7 +173,7 @@ impl<R: MmapBytesReader> IpcReader<R> {
 
         let reader = read::FileReader::new(self.reader, metadata, self.projection, self.n_rows);
 
-        finish_reader(reader, rechunk, None, predicate, &schema, self.row_count)
+        finish_reader(reader, rechunk, None, predicate, &schema, self.row_index)
     }
 }
 
@@ -194,7 +194,7 @@ impl<R: MmapBytesReader> SerReader<R> for IpcReader<R> {
             n_rows: None,
             columns: None,
             projection: None,
-            row_count: None,
+            row_index: None,
             memmap: true,
             metadata: None,
             schema: None,
@@ -230,6 +230,6 @@ impl<R: MmapBytesReader> SerReader<R> for IpcReader<R> {
 
         let ipc_reader =
             read::FileReader::new(self.reader, metadata.clone(), self.projection, self.n_rows);
-        finish_reader(ipc_reader, rechunk, None, None, &schema, self.row_count)
+        finish_reader(ipc_reader, rechunk, None, None, &schema, self.row_index)
     }
 }
diff --git a/crates/polars-io/src/ipc/ipc_stream.rs b/crates/polars-io/src/ipc/ipc_stream.rs
@@ -69,7 +69,7 @@ pub struct IpcStreamReader<R> {
     n_rows: Option<usize>,
     projection: Option<Vec<usize>>,
     columns: Option<Vec<String>>,
-    row_count: Option<RowCount>,
+    row_index: Option<RowIndex>,
     metadata: Option<StreamMetadata>,
 }
 
@@ -95,9 +95,9 @@ impl<R: Read> IpcStreamReader<R> {
         self
     }
 
-    /// Add a `row_count` column.
-    pub fn with_row_count(mut self, row_count: Option<RowCount>) -> Self {
-        self.row_count = row_count;
+    /// Add a row index column.
+    pub fn with_row_index(mut self, row_index: Option<RowIndex>) -> Self {
+        self.row_index = row_index;
         self
     }
 
@@ -146,7 +146,7 @@ where
             n_rows: None,
             columns: None,
             projection: None,
-            row_count: None,
+            row_index: None,
             metadata: None,
         }
     }
@@ -177,7 +177,7 @@ where
             metadata.schema.clone()
         };
 
-        let include_row_count = self.row_count.is_some();
+        let include_row_index = self.row_index.is_some();
         let ipc_reader =
             read::StreamReader::new(&mut self.reader, metadata.clone(), sorted_projection);
         finish_reader(
@@ -186,23 +186,27 @@ where
             self.n_rows,
             None,
             &schema,
-            self.row_count,
+            self.row_index,
         )
-        .map(|df| fix_column_order(df, self.projection, include_row_count))
+        .map(|df| fix_column_order(df, self.projection, include_row_index))
     }
 }
 
-fn fix_column_order(df: DataFrame, projection: Option<Vec<usize>>, row_count: bool) -> DataFrame {
+fn fix_column_order(
+    df: DataFrame,
+    projection: Option<Vec<usize>>,
+    include_row_index: bool,
+) -> DataFrame {
     if let Some(proj) = projection {
-        let offset = usize::from(row_count);
+        let offset = usize::from(include_row_index);
         let mut args = (0..proj.len()).zip(proj).collect::<Vec<_>>();
         // first el of tuple is argument index
         // second el is the projection index
         args.sort_unstable_by_key(|tpl| tpl.1);
         let cols = df.get_columns();
 
         let iter = args.iter().map(|tpl| cols[tpl.0 + offset].clone());
-        let cols = if row_count {
+        let cols = if include_row_index {
             let mut new_cols = vec![df.get_columns()[0].clone()];
             new_cols.extend(iter);
             new_cols
@@ -242,7 +246,7 @@ pub struct IpcStreamWriter<W> {
 
 use polars_core::frame::ArrowChunk;
 
-use crate::RowCount;
+use crate::RowIndex;
 
 impl<W> IpcStreamWriter<W> {
     /// Set the compression used. Defaults to None.

diff --git a/crates/polars-io/src/ipc/mmap.rs b/crates/polars-io/src/ipc/mmap.rs
@@ -99,7 +99,7 @@ impl<R: MmapBytesReader> IpcReader<R> {
                     self.n_rows,
                     predicate,
                     &schema,
-                    self.row_count.clone(),
+                    self.row_index.clone(),
                 )
             },
             None => polars_bail!(ComputeError: "cannot memory-map, you must provide a file"),

diff --git a/crates/polars-io/src/lib.rs b/crates/polars-io/src/lib.rs
@@ -94,7 +94,7 @@ pub(crate) fn finish_reader<R: ArrowReader>(
     n_rows: Option<usize>,
     predicate: Option<Arc<dyn PhysicalIoExpr>>,
     arrow_schema: &ArrowSchema,
-    row_count: Option<RowCount>,
+    row_index: Option<RowIndex>,
 ) -> PolarsResult<DataFrame> {
     use polars_core::utils::accumulate_dataframes_vertical;
 
@@ -106,8 +106,8 @@ pub(crate) fn finish_reader<R: ArrowReader>(
         num_rows += batch.len();
         let mut df = DataFrame::try_from((batch, arrow_schema.fields.as_slice()))?;
 
-        if let Some(rc) = &row_count {
-            df.with_row_count_mut(&rc.name, Some(current_num_rows + rc.offset));
+        if let Some(rc) = &row_index {
+            df.with_row_index_mut(&rc.name, Some(current_num_rows + rc.offset));
         }
 
         if let Some(predicate) = &predicate {