let m = PyModule::new(py, "file")?;

parent.add_submodule(&m)?;

install_module("vortex._lib.file", &m)?;

m.add_function(wrap_pyfunction!(open, &m)?)?;

m.add_function(wrap_pyfunction!(_reopen, &m)?)?;

m.add_class::<PyVortexFile>()?;

Ok(())

py: Python,

path: &str,

store: Option<AnyVortexStore>,

without_segment_cache: bool,

) -> PyVortexResult<PyVortexFile> {

let had_store = store.is_some();

let owned_path = path.to_string();

let vxf = py.detach(move || {

current_runtime().block_on(async move {

let mut options = session().open_options();

if !without_segment_cache {

// TODO(ngates): use a globally shared segment cache for all files

options = options.with_segment_cache(Arc::new(MokaSegmentCache::new(256 << 20)));

}

match resolve_store(path, store.map(|x| x.into_inner()))? {

ResolvedStore::ObjectStore(store, path) => {

options.open_object_store(&store, path.as_ref()).await

}

ResolvedStore::Path(path) => options.open_path(path).await,

}

})

})?;

Ok(PyVortexFile {

vxf,

path: owned_path,

had_store,

without_segment_cache,

})

vxf: VortexFile,

/// The path this file was opened from, retained so that it can be reopened in another process.

path: String,

/// Whether an explicit object store was passed to [`open`]. Object stores are not picklable, so

/// such a file cannot be reopened by path alone.

had_store: bool,

without_segment_cache: bool,

fn __len__(slf: PyRef<Self>) -> PyResult<usize> {

Ok(usize::try_from(slf.vxf.row_count())?)

}

/// The path or URL this file was opened from.

///

/// Returns

/// -------

/// :class:`.str`

#[getter]

fn path(slf: PyRef<Self>) -> String {

slf.path.clone()

}

/// Support for Python's pickle protocol: the file is reopened by path in the receiving process.

///

/// Only the path is transferred, never any read state or segment cache, so this is cheap enough

/// to send a file to every worker of a multiprocessing pool or Ray job.

///

/// Raises

/// ------

/// :class:`TypeError`

/// If the file was opened with an explicit ``store``, since object stores cannot be

/// pickled. Pass the URL to :func:`vortex.open` in the worker instead.

fn __reduce__<'py>(slf: PyRef<'py, Self>) -> PyResult<(Bound<'py, PyAny>, (String, bool))> {

if slf.had_store {

return Err(PyTypeError::new_err(

"cannot pickle a VortexFile opened with an explicit store, because object stores \

));

}

let py = slf.py();

let module = PyModule::import(py, "vortex._lib.file")?;

let reopen = module.getattr(intern!(py, "_reopen"))?;

Ok((reopen, (slf.path.clone(), slf.without_segment_cache)))

}

#[getter]

fn dtype(slf: Bound<Self>) -> PyResult<Bound<PyDType>> {

PyDType::init(slf.py(), slf.get().vxf.dtype().clone())

}

#[pyo3(signature = (projection = None, *, expr = None, limit = None, indices = None, batch_size = None))]

fn scan(

slf: Bound<Self>,

projection: Option<PyIntoProjection>,

expr: Option<PyExpr>,

limit: Option<u64>,

indices: Option<PyArrayRef>,

batch_size: Option<usize>,

) -> PyVortexResult<PyArrayIterator> {

let vxf = slf.get().vxf.clone();

let projection = projection.map(|p| p.0);

let expr = expr.map(|e| e.into_inner());

let indices = indices.map(|i| i.into_inner());

slf.py().detach(move || {

let session = session();

let mut ctx = session.create_execution_ctx();

let builder =

scan_builder(&vxf, projection, expr, limit, indices, batch_size, &mut ctx)?;

let runtime = current_runtime();

Ok(PyArrayIterator::new(Box::new(

builder.into_array_iter(&runtime)?,

)))

})

}

#[pyo3(signature = (projection = None, *, expr = None, limit = None, indices = None, batch_size = None))]

fn prepare(

slf: Bound<Self>,

projection: Option<PyIntoProjection>,

expr: Option<PyExpr>,

limit: Option<u64>,

indices: Option<PyArrayRef>,

batch_size: Option<usize>,

) -> PyVortexResult<PyRepeatedScan> {

let vxf = slf.get().vxf.clone();

let projection = projection.map(|p| p.0);

let expr = expr.map(|e| e.into_inner());

let indices = indices.map(|i| i.into_inner());

let scan = slf.py().detach(move || {

let session = session();

let mut ctx = session.create_execution_ctx();

scan_builder(&vxf, projection, expr, limit, indices, batch_size, &mut ctx)?.prepare()

})?;

Ok(PyRepeatedScan {

scan: Arc::new(scan),

row_count: slf.get().vxf.row_count(),

})

}

#[pyo3(signature = (projection = None, *, expr = None, limit = None, batch_size = None, schema = None))]

fn to_arrow(

slf: Bound<Self>,

projection: Option<PyIntoProjection>,

expr: Option<PyExpr>,

limit: Option<u64>,

batch_size: Option<usize>,

schema: Option<&Bound<PyAny>>,

) -> PyVortexResult<Py<PyAny>> {

let vxf = slf.get().vxf.clone();

let schema = schema

.map(|schema| Schema::from_pyarrow(&schema.as_borrowed()))

.transpose()?

.map(Arc::new);

let runtime = current_runtime();

let reader = slf.py().detach(|| {

let filter = expr

.map(|e| {

e.into_inner()

.optimize_recursive(vxf.dtype())?

.bind(vxf.dtype())

})

.transpose()?;

let projection = projection

.map(|p| p.0)

.unwrap_or_else(root)

.optimize_recursive(vxf.dtype())?

.bind(vxf.dtype())?;

let mut builder = vxf

.scan()?

.with_some_filter(filter)

.with_projection(projection);

if let Some(limit) = limit {

builder = builder.with_limit(limit);

}

if let Some(batch_size) = batch_size {

builder = builder.with_split_by(SplitBy::RowCount(batch_size));

}

let schema = match schema {

Some(schema) => schema,

None => Arc::new(session().arrow().to_arrow_schema(&builder.dtype()?)?),

};

builder.into_record_batch_reader(schema, &runtime)

})?;

let rbr: Box<dyn RecordBatchReader + Send> = Box::new(reader);

Ok(rbr.into_pyarrow(slf.py())?)

}

fn to_dataset(slf: Bound<Self>) -> PyVortexResult<PyVortexDataset> {

Ok(PyVortexDataset::try_new(slf.get().vxf.clone())?)

}

#[pyo3(signature = (*))]

pub fn splits(&self) -> PyVortexResult<Vec<(u64, u64)>> {

Ok(self

.vxf

.splits()?

.into_iter()

.map(|x| (x.start, x.end))

.collect())

}

vxf: &VortexFile,

projection: Option<Expression>,

expr: Option<Expression>,

limit: Option<u64>,

indices: Option<ArrayRef>,

batch_size: Option<usize>,

ctx: &mut ExecutionCtx,

) -> VortexResult<ScanBuilder<ArrayRef>> {

let projection = projection

.unwrap_or_else(root)

.optimize_recursive(vxf.dtype())?

.bind(vxf.dtype())?;

let expr = expr

.map(|expr| expr.optimize_recursive(vxf.dtype())?.bind(vxf.dtype()))

.transpose()?;

let mut builder = vxf

.scan()?

.with_some_filter(expr)

.with_projection(projection);

if let Some(limit) = limit {

builder = builder.with_limit(limit);

}

if let Some(indices) = indices {

let casted = indices.cast(DType::Primitive(PType::U64, NonNullable))?;

let indices = casted.execute::<PrimitiveArray>(ctx)?.into_buffer::<u64>();

builder = builder.with_row_indices(StrictSortedBuffer::try_new(indices)?);

}

if let Some(batch_size) = batch_size {

builder = builder.with_split_by(SplitBy::RowCount(batch_size));

}

Ok(builder)

type Error = PyErr;

fn extract(ob: Borrowed<'_, 'py, PyAny>) -> Result<Self, Self::Error> {

// If it's a list of strings, convert to a column selection.

if let Ok(py_list) = ob.cast::<PyList>() {

let cols = py_list

.iter()

.map(|item| item.extract::<String>())

.collect::<PyResult<Vec<String>>>()?;

return Ok(PyIntoProjection(select(

cols.into_iter().collect::<FieldNames>(),

root(),

)));

}

// If it's an expression, just return it.

if let Ok(py_expr) = ob.cast::<PyExpr>() {

return Ok(PyIntoProjection(py_expr.get().inner().clone()));

}

Err(PyTypeError::new_err(

"projection must be a list of strings or a vortex.Expr",

))

}