refactor(rust): Request for preliminary directional check-in: making file loading much better factored #12750

itamarst · 2023-11-28T16:31:50Z

Motivation

User stories

Users want to be able to load CSVs from file-like Python objects (Support BytesIO, StringIO etc. in scan_csv() #12617)
Likewise for Parquet (scan_parquet from io.BytesIO() #10413)

Issues with the current code

Structural issues

In polars-lazy, the various lazy file readers had a bad pattern where there was both a path and paths on the structs.
- The path had two completely different meanings: a glob, and a specific file. A glob would be converted to a specific file with a specific path.
- The paths would also be converted into specific files with a specific path.
- In short, this is an implicit state machine with ... 6 states? Some of which were never addressed.
In addition, any time time a PathBuf was used for a specific path to a file or as a glob, it had two potential meanings: a local path on the filesystem, or a cloud URI. Dispatch was done based on is_cloud_url().

Both these cases ought to be modeled via the type system, not string inspection or implicit, vague state machines.

Functionality issues

Because all loading was assumed to be from a PathBuf, loading from a Python file-like object is not possible.

Proposed design changes

Switch from path + paths in LazyFileListReader to a single SpecificOrMultipleReaderFactories. This is still not great, since there's still a requirement for the code to do things in a certain order if it doesn't want a panic; a better refactor would involve splitting the classes into a builder phase and a final phase. But it's a large enough branch as it is, and at least now there's only two states instead of six!
Use of PathBuf for things to load is replaced with a ReaderFactory; for now it's either a filesystem path or a URI, but a Python factory function that returns file-like objects could be added. It's a factory to allow for opening a file multiple times, and because of the need for (de)serialization.

mmap.rs is probably the wrong place for ReaderFactory, it should probably be elsewhere.

I did some of the refactoring, but there's a lot more, and it will percolate throughout the code base, so this seemed like a good place to get some feedback.

State of the PR

There's lots and lots of code that hasn't been convered from PathBuf to ReaderFactory, and I haven't quite proven to myself that it will support python file-like objects as well. But it probably will?

itamarst · 2023-11-28T16:42:11Z

crates/polars-io/src/mmap.rs

@@ -84,3 +91,41 @@ impl<'a, T: 'a + MmapBytesReader> From<&'a T> for ReaderBytes<'a> {
        }
    }
 }
+
+/// Create MmapBytesReaders for a specific "file", either locally or remotely.


This replaces the use of PathBuf as a source of data to be read.

itamarst · 2023-11-28T16:44:34Z

crates/polars-lazy/src/scan/csv.rs


 #[derive(Clone)]
 #[cfg(feature = "csv")]
 pub struct LazyCsvReader<'a> {
-    path: PathBuf,


Notice there are multiple potential states in the original code.

Expected states:

path is glob, paths is empty.

path is specific path, paths is empty.

path is empty, paths is non-empty list of specific paths.

Unexpected and hopefully not present states:

path is glob, paths is non-empty list of specific paths.

path is specific path, paths is non-empty list of specific paths.

path is empty, paths is empty.

ritchie46 · 2023-11-28T16:50:23Z

I hope to get to this in a few days. A bit busy atm.

crates/polars-lazy/src/scan/file_list_reader.rs

crates/polars-lazy/src/scan/ndjson.rs

crates/polars-lazy/src/scan/ipc.rs

itamarst · 2023-11-28T16:59:06Z

crates/polars-plan/src/logical_plan/builder.rs

@@ -199,7 +210,10 @@ impl LogicalPlanBuilder {
        // We set the hive partitions of the first path to determine the schema.
        // On iteration the partition values will be re-set per file.
        if hive_partitioning {
-            file_info.init_hive_partitions(path.as_path())?;
+            if let Some(uri) = uri {


I think this is correct and hive only applies to cloud URLs? This is why it's better not to have "a string" as a datastructure, and explicitly separate local paths from cloud URLs in the type system.

crates/polars-lazy/src/scan/csv.rs

crates/polars-io/src/mmap.rs

crates/polars-lazy/src/scan/file_list_reader.rs

itamarst · 2023-12-05T17:44:07Z

Thank you, I will try to update the sketch later this week.

itamarst · 2023-12-06T18:40:39Z

I hadn't quite realized before the ways in which this is exposed as part of the public Rust API, mostly I was thinking about it from perspective of implementation. So now I'm thinking about that half more in my sketch.

…lity methods.

itamarst · 2023-12-06T18:57:34Z

OK, here's a new design. It turns LazyFileListReader into a single-state factory, rather than the 6 states in main or the 2 states in previous PR iteration. It also makes more explicit what happens when you load e.g. a string (personally I am uncomfortable with idea that globbing might be used by mistake so I think it's helpful to have a few extra more-explicit utility methods for loading).

itamarst · 2023-12-06T18:58:48Z

crates/polars-lazy/src/scan/ipc.rs

@@ -44,9 +44,8 @@ impl LazyIpcReader {
 }

 impl LazyFileListReader for LazyIpcReader {
-    fn finish_no_glob(self) -> PolarsResult<LazyFrame> {
+    fn finish_no_glob(self, path: &Path) -> PolarsResult<LazyFrame> {


This and similar changes should actually be load_specific(self, location: ScanLocation).

itamarst · 2023-12-06T19:04:06Z

crates/polars/src/docs/lazy.rs

@@ -112,10 +112,10 @@
 //! use polars::prelude::*;
 //! # fn example() -> PolarsResult<()> {
 //!
-//!  let df = LazyCsvReader::new("reddit.csv")
+//!  let df = LazyCsvReader::new()


This diff shows how the public API differs. By moving the path/paths/glob/etc to be loaded to the load_* methods (which consume the factory and finish the LazyFrame construction), there is no longer a requirement for weird internal state tracking as there was previously with path/paths.

itamarst · 2023-12-06T19:12:46Z

crates/polars-lazy/src/scan/csv.rs

@@ -233,20 +228,12 @@ impl<'a> LazyCsvReader<'a> {
    /// Modify a schema before we run the lazy scanning.
    ///
    /// Important! Run this function latest in the builder!
-    pub fn with_schema_modify<F>(self, f: F) -> PolarsResult<Self>
+    pub fn with_schema_modify<F>(mut self, f: F, location: &ScanLocation) -> PolarsResult<Self>


This is kinda broken at this stage in the refactor, but insofar as it's only use in Polars is in the py-polars wrapper I assume it can be made to work with some further tweaking.

itamarst · 2023-12-18T16:02:17Z

Sounds like Polars is going to stabilize APIs soon? If so, would love to get this in first since it will become much harder afterwards.

stinodego · 2024-02-14T13:19:11Z

Sounds like Polars is going to stabilize APIs soon? If so, would love to get this in first since it will become much harder afterwards.

We do not consider the Rust API stable yet and are not planning to soon. Breaking changes are allowed in every PR.

The Python versions are going to 1.0.0 soon though, but that should not impact this PR.

Unfortunately, I will not be much help in reviewing this PR as this is not my expertise.

pythonspeed added 5 commits November 22, 2023 17:18

Sketch of better set of interfaces

9f56842

More refactoring, work in progress

83d07cb

Continue refactoring to use ReaderFactory instead of strings.

e748902

Some docs

7d7c4a1

More refactoring to use ReaderFactory

43a3b11

itamarst requested review from ritchie46 and orlp as code owners November 28, 2023 16:31

itamarst marked this pull request as draft November 28, 2023 16:32

itamarst commented Nov 28, 2023

View reviewed changes

Actually use ReaderFactory

86759a5

pythonspeed added 2 commits November 28, 2023 11:52

Better names

e1a5eae

Correct type

390f7a2

itamarst commented Nov 28, 2023

View reviewed changes

Maybe clearer about what's going on.

fdaed9e

ritchie46 reviewed Dec 2, 2023

View reviewed changes

crates/polars-lazy/src/scan/csv.rs Outdated Show resolved Hide resolved

crates/polars-io/src/mmap.rs Outdated Show resolved Hide resolved

crates/polars-lazy/src/scan/file_list_reader.rs Outdated Show resolved Hide resolved

pythonspeed added 2 commits December 6, 2023 13:54

Make the LazyFileListReader both more explicit and requiring less uti…

ad83865

…lity methods.

Better naming.

bb14ecb

itamarst commented Dec 6, 2023

View reviewed changes

ReaderFactory -> ScanLocation

7a0406c

itamarst commented Dec 6, 2023

View reviewed changes

pythonspeed added 2 commits December 6, 2023 14:08

Get rid of more references to ReaderFactory.

45d3fe0

Make it kinda usable.

fe1f08e

itamarst commented Dec 6, 2023

View reviewed changes

itamarst requested a review from ritchie46 December 6, 2023 19:13

stinodego changed the title ~~Request for preliminary directional check-in: making file loading much better factored~~ refactor(rust): Request for preliminary directional check-in: making file loading much better factored Feb 26, 2024

github-actions bot added internal An internal refactor or improvement rust Related to Rust Polars labels Feb 26, 2024

ritchie46 force-pushed the main branch from 0a696ff to 9c29683 Compare July 28, 2024 08:11

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

refactor(rust): Request for preliminary directional check-in: making file loading much better factored #12750

refactor(rust): Request for preliminary directional check-in: making file loading much better factored #12750

itamarst commented Nov 28, 2023 •

edited

Loading

itamarst Nov 28, 2023

itamarst Nov 28, 2023 •

edited

Loading

ritchie46 commented Nov 28, 2023

itamarst Nov 28, 2023

itamarst commented Dec 5, 2023

itamarst commented Dec 6, 2023

itamarst commented Dec 6, 2023 •

edited

Loading

itamarst Dec 6, 2023 •

edited

Loading

itamarst Dec 6, 2023 •

edited

Loading

itamarst Dec 6, 2023 •

edited

Loading

itamarst commented Dec 18, 2023

stinodego commented Feb 14, 2024 •

edited

Loading

refactor(rust): Request for preliminary directional check-in: making file loading much better factored #12750

Are you sure you want to change the base?

refactor(rust): Request for preliminary directional check-in: making file loading much better factored #12750

Conversation

itamarst commented Nov 28, 2023 • edited Loading

Motivation

User stories

Issues with the current code

Structural issues

Functionality issues

Proposed design changes

State of the PR

itamarst Nov 28, 2023

Choose a reason for hiding this comment

itamarst Nov 28, 2023 • edited Loading

Choose a reason for hiding this comment

ritchie46 commented Nov 28, 2023

itamarst Nov 28, 2023

Choose a reason for hiding this comment

itamarst commented Dec 5, 2023

itamarst commented Dec 6, 2023

itamarst commented Dec 6, 2023 • edited Loading

itamarst Dec 6, 2023 • edited Loading

Choose a reason for hiding this comment

itamarst Dec 6, 2023 • edited Loading

Choose a reason for hiding this comment

itamarst Dec 6, 2023 • edited Loading

Choose a reason for hiding this comment

itamarst commented Dec 18, 2023

stinodego commented Feb 14, 2024 • edited Loading

itamarst commented Nov 28, 2023 •

edited

Loading

itamarst Nov 28, 2023 •

edited

Loading

itamarst commented Dec 6, 2023 •

edited

Loading

itamarst Dec 6, 2023 •

edited

Loading

itamarst Dec 6, 2023 •

edited

Loading

itamarst Dec 6, 2023 •

edited

Loading

stinodego commented Feb 14, 2024 •

edited

Loading