-
Notifications
You must be signed in to change notification settings - Fork 199
Expand file tree
/
Copy pathscan.rs
More file actions
95 lines (82 loc) · 2.88 KB
/
Copy pathscan.rs
File metadata and controls
95 lines (82 loc) · 2.88 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors
use std::sync::Arc;
use pyo3::exceptions::PyIndexError;
use pyo3::prelude::*;
use vortex::array::ArrayRef;
use vortex::array::VortexSessionExecute;
use vortex::error::VortexResult;
use vortex::layout::scan::repeated_scan::RepeatedScan;
use vortex::scalar::Scalar;
use crate::current_runtime;
use crate::error::PyVortexResult;
use crate::install_module;
use crate::iter::PyArrayIterator;
use crate::scalar::PyScalar;
use crate::session::session;
pub(crate) fn init(py: Python, parent: &Bound<PyModule>) -> PyResult<()> {
let m = PyModule::new(py, "scan")?;
parent.add_submodule(&m)?;
install_module("vortex._lib.scan", &m)?;
m.add_class::<PyRepeatedScan>()?;
Ok(())
}
#[pyclass(name = "RepeatedScan", module = "vortex", frozen)]
pub struct PyRepeatedScan {
pub scan: Arc<RepeatedScan<ArrayRef>>,
pub row_count: u64,
}
#[pymethods]
impl PyRepeatedScan {
#[pyo3(signature = (*, start = None, stop = None))]
fn execute(
slf: Bound<Self>,
start: Option<u64>,
stop: Option<u64>,
) -> PyVortexResult<PyArrayIterator> {
let row_count = slf.get().row_count;
let row_range = match (start, stop) {
(Some(start), Some(stop)) => Some(start..stop),
(Some(start), None) => Some(start..row_count),
(None, Some(stop)) => Some(0..stop),
(None, None) => None,
};
let scan = Arc::clone(&slf.get().scan);
slf.py().detach(move || {
let runtime = current_runtime();
Ok(PyArrayIterator::new(Box::new(
scan.execute_array_iter(row_range, &runtime)?,
)))
})
}
fn scalar_at(slf: Bound<Self>, index: u64) -> PyVortexResult<Bound<PyScalar>> {
let row_count = slf.get().row_count;
if index >= row_count {
return Err(PyIndexError::new_err(format!(
"Index out of bounds: {} >= {}",
index, row_count
))
.into());
}
let scan = Arc::clone(&slf.get().scan);
let scalar = slf.py().detach(move || -> VortexResult<Option<Scalar>> {
let session = session();
let runtime = current_runtime();
for batch in scan.execute_array_iter(Some(index..index + 1), &runtime)? {
let array = batch?;
if array.is_empty() {
continue;
}
let scalar = array.execute_scalar(0, &mut session.create_execution_ctx())?;
return Ok(Some(scalar));
}
Ok(None)
})?;
match scalar {
Some(scalar) => Ok(PyScalar::init(slf.py(), scalar)?),
None => {
Err(PyIndexError::new_err(format!("Index {} not found in the scan", index)).into())
}
}
}
}