PyO3 / PyO3/pyo3

Integrating Polars with Datafusion

Open
#4,843 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Rust
Stars
16.2k
Forks
1k
Avg merge
2d 6h
Merged PRs (30d)
66

Description

Hey I'm working on implementing streaming functionality for Polars on top of some extensions to DataFusion.
For this purpose I added streaming support for AnonymousScan in Polars:

pub struct RangeOperationScan {
    pub(crate) df_iter: Arc<Mutex<SendableRecordBatchStream>>,
}


impl AnonymousScan for SomeOperationScan {
    fn as_any(&self) -> &dyn std::any::Any {
        self
    }

    fn scan(&self, scan_opts: AnonymousScanArgs) -> PolarsResult<polars::prelude::DataFrame> {
        !todo!("Only streaming is supported")
    }

    fn next_batch(
        &self,
        scan_opts: AnonymousScanArgs,
    ) -> PolarsResult<Option<polars::prelude::DataFrame>> {
        let mutex_stream = Arc::clone(&self.df_iter);
        thread::spawn(move ||{
            let rt = Runtime::new().unwrap();
            let result = rt.block_on(mutex_stream.lock().unwrap().next()); // <-- I think here is the problem
            match result {
                Some(batch) => {
                    let rb = batch.unwrap();
                    let schema_polars = convert_arrow_rb_schema_to_polars_df_schema(&rb.schema())?; 
                    let df = convert_arrow_rb_to_polars_df(&rb, &schema_polars)?;
                    Ok(Some(df))
                },
                None => Ok(None),
            }
        }).join().unwrap()
        

an then:

fn lazy_range_operation_scan(
    py: Python<'_>,
    py_ctx: &PyBioSessionContext,
    df_path1: String,
    df_path2: String,
    range_options: RangeOptions,
) -> PyResult<PyLazyFrame> {
    py.allow_threads(|| {
        
// some code removed
      
        let rt = Runtime::new().unwrap();
        let ctx = &py_ctx.ctx;

        let args = ScanArgsAnonymous {
            schema: Some(Arc::new()),
            ..ScanArgsAnonymous::default()
        };
      
   // some code removed 
        let stream = rt.block_on(df?.execute_stream())?;
        let scan = RangeOperationScan {
            df_iter: Arc::new(Mutex::new(stream)),
        };
        let function = Arc::new(scan);
        let lf = LazyFrame::anonymous_scan(function, args).unwrap();
        Ok(lf.into())
}

Everything works like charm when there is only Datafusion runs using a single-thread (or multithread with Polars Rust API). If i try to run in multi-threaded mode in Python I suspect that main thread quits causing Datafusion tasks get cancelled:

thread '<unnamed>' panicked at src/scan.rs:97:36:
called `Result::unwrap()` on an `Err` value: External(Internal("Non Panic Task error: task 43 was cancelled"))
stack backtrace:
   0:        0x34c592498 - <unknown>
   1:        0x34c5b4834 - <unknown>

not sure if it's possible to prevent that five that next_batch method is called externally by the Polars engine.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

The report centers on src/scan.rs:97 and the AnonymousScan::next_batch implementation, with lazy_range_operation_scan creating the DataFusion stream. Reproduce the Python multi-threaded failure and trace the lifetime of the runtime and execute_stream task; done means establishing whether the cancellation can be prevented or documenting the binding limitation.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, rust
Domain
backend
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.