Dataset API (#5949)

# Dataset API

A unified API to sit between query engines and data sources that
preserves support for late materialization, deferred decompression, and
alternate device buffers.

# Open Questions

* What does capabilities negotiation look like? e.g. for unsupported
encodings?
* What about shared dictionaries?
* Is deserialize_split in the right place? Without serde, splits can
share state in the DataSourceScan. With serde, there's no where for
shared state to live. Perhaps we should reconstruct a datasource on the
worker? Should the datasource be serializable?

---------

Signed-off-by: Nicholas Gates <[email protected]>
Co-authored-by: Dmitrii Blaginin <[email protected]>

Author: gatesn

Cargo.lock

@@ -29,6 +29,7 @@ vortex-mask = { workspace = true }
 vortex-metrics = { workspace = true }
 vortex-session = { workspace = true }
 
+async-trait = { workspace = true }
 bit-vec = { workspace = true }
 futures = { workspace = true }
 itertools = { workspace = true }

vortex-scan/Cargo.toml

@@ -0,0 +1,135 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+#![deny(missing_docs)]
+
+//! The Vortex Scan API implements an abstract table scan interface that can be used to
+//! read data from various data sources.
+//!
+//! It supports arbitrary projection expressions, filter expressions, and limit pushdown as well
+//! as mechanisms for parallel and distributed execution via splits.
+//!
+//! The API is currently under development and may change in future releases, however we hope to
+//! stabilize into stable C ABI for use within foreign language bindings.
+//!
+//! ## Open Issues
+//!
+//! * We probably want to make the DataSource serializable as well, so that we can share
+//!   source-level state with workers, separately from split serialization.
+//! * We should add a way for the client to negotiate capabilities with the data source, for
+//!   example which encodings it knows about.
+
+use std::any::Any;
+use std::sync::Arc;
+
+use async_trait::async_trait;
+use futures::stream::BoxStream;
+use vortex_array::expr::Expression;
+use vortex_array::stream::SendableArrayStream;
+use vortex_dtype::DType;
+use vortex_error::VortexResult;
+use vortex_session::VortexSession;
+
+/// Create a Vortex source from serialized configuration.
+///
+/// Providers can be registered with Vortex under a specific
+#[async_trait(?Send)]
+pub trait DataSourceProvider: 'static {
+    /// Attempt to initialize a new source.
+    ///
+    /// Returns `Ok(None)` if the provider cannot handle the given URI.
+    async fn initialize(
+        &self,
+        uri: String,
+        session: &VortexSession,
+    ) -> VortexResult<Option<DataSourceRef>>;
+}
+
+/// A reference-counted data source.
+pub type DataSourceRef = Arc<dyn DataSource>;
+
+/// A data source represents a streamable dataset that can be scanned with projection and filter
+/// expressions. Each scan produces splits that can be executed (potentially in parallel) to read
+/// data. Each split can be serialized for remote execution.
+///
+/// The DataSource may be used multiple times to create multiple scans, whereas each scan and each
+/// split of a scan can only be consumed once.
+#[async_trait]
+pub trait DataSource: 'static + Send + Sync {
+    /// Returns the dtype of the source.
+    fn dtype(&self) -> &DType;
+
+    /// Returns an estimate of the row count of the source.
+    fn row_count_estimate(&self) -> Estimate<u64>;
+
+    /// Returns a scan over the source.
+    async fn scan(&self, scan_request: ScanRequest) -> VortexResult<DataSourceScanRef>;
+
+    /// Serialize a split from this data source.
+    fn serialize_split(&self, split: &dyn Split) -> VortexResult<Vec<u8>>;
+
+    /// Deserialize a split that was previously serialized from a compatible data source.
+    fn deserialize_split(&self, data: &[u8]) -> VortexResult<SplitRef>;
+}
+
+/// A request to scan a data source.
+#[derive(Debug, Clone, Default)]
+pub struct ScanRequest {
+    /// Projection expression, `None` implies `root()`.
+    pub projection: Option<Expression>,
+    /// Filter expression, `None` implies no filter.
+    pub filter: Option<Expression>,
+    /// Optional limit on the number of rows to scan.
+    pub limit: Option<u64>,
+}
+
+/// A boxed data source scan.
+pub type DataSourceScanRef = Box<dyn DataSourceScan>;
+
+/// A data source scan produces splits that can be executed to read data from the source.
+#[async_trait]
+pub trait DataSourceScan: 'static + Send {
+    /// The returned dtype of the scan.
+    fn dtype(&self) -> &DType;
+
+    /// An estimate of the remaining splits.
+    fn remaining_splits_estimate(&self) -> Estimate<usize>;
+
+    /// Returns the next batch of splits to be processed.
+    ///
+    /// This should not return _more_ than `max_splits` splits, but may return fewer.
+    async fn next_splits(&mut self, max_splits: usize) -> VortexResult<Vec<SplitRef>>;
+}
+
+/// A stream of splits.
+pub type SplitStream = BoxStream<'static, VortexResult<SplitRef>>;
+
+/// A reference-counted split.
+pub type SplitRef = Box<dyn Split>;
+
+/// A split represents a unit of work that can be executed to produce a stream of arrays.
+pub trait Split: 'static + Send {
+    /// Downcast the split to a concrete type.
+    fn as_any(&self) -> &dyn Any;
+
+    /// Executes the split.
+    fn execute(self: Box<Self>) -> VortexResult<SendableArrayStream>;
+
+    /// Returns an estimate of the row count for this split.
+    fn row_count_estimate(&self) -> Estimate<u64>;
+
+    /// Returns an estimate of the byte size for this split.
+    fn byte_size_estimate(&self) -> Estimate<u64>;
+}
+
+/// An estimate that can be exact, an upper bound, or unknown.
+#[derive(Default)]
+pub enum Estimate<T> {
+    /// The exact value.
+    Exact(T),
+    /// An upper bound on the value.
+    UpperBound(T),
+    /// The value is unknown.
+    #[default]
+    Unknown,
+}

vortex-scan/src/api.rs

@@ -0,0 +1,148 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::any::Any;
+use std::collections::VecDeque;
+
+use async_trait::async_trait;
+use futures::StreamExt;
+use futures::future::BoxFuture;
+use futures::stream;
+use vortex_array::ArrayRef;
+use vortex_array::stream::ArrayStreamAdapter;
+use vortex_array::stream::SendableArrayStream;
+use vortex_dtype::DType;
+use vortex_error::VortexExpect;
+use vortex_error::VortexResult;
+use vortex_error::vortex_bail;
+use vortex_layout::LayoutReaderRef;
+use vortex_session::VortexSession;
+
+use crate::ScanBuilder;
+use crate::api::DataSource;
+use crate::api::DataSourceScan;
+use crate::api::DataSourceScanRef;
+use crate::api::Estimate;
+use crate::api::ScanRequest;
+use crate::api::Split;
+use crate::api::SplitRef;
+
+/// An implementation of a [`DataSource`] that reads data from a [`LayoutReaderRef`].
+pub struct LayoutReaderDataSource {
+    reader: LayoutReaderRef,
+    session: VortexSession,
+}
+
+impl LayoutReaderDataSource {
+    /// Creates a new [`LayoutReaderDataSource`].
+    pub fn new(reader: LayoutReaderRef, session: VortexSession) -> Self {
+        Self { reader, session }
+    }
+}
+
+#[async_trait]
+impl DataSource for LayoutReaderDataSource {
+    fn dtype(&self) -> &DType {
+        self.reader.dtype()
+    }
+
+    fn row_count_estimate(&self) -> Estimate<u64> {
+        Estimate::Exact(self.reader.row_count())
+    }
+
+    async fn scan(&self, scan_request: ScanRequest) -> VortexResult<DataSourceScanRef> {
+        let mut builder = ScanBuilder::new(self.session.clone(), self.reader.clone());
+
+        if let Some(projection) = scan_request.projection {
+            builder = builder.with_projection(projection);
+        }
+
+        if let Some(filter) = scan_request.filter {
+            builder = builder.with_filter(filter);
+        }
+
+        if let Some(limit) = scan_request.limit {
+            // TODO(ngates): ScanBuilder limit should be u64
+            let limit = usize::try_from(limit)?;
+            builder = builder.with_limit(limit);
+        }
+
+        let scan = builder.prepare()?;
+        let dtype = scan.dtype().clone();
+        let splits = scan.execute(None)?;
+
+        Ok(Box::new(LayoutReaderScan {
+            dtype,
+            splits: VecDeque::from_iter(splits),
+        }))
+    }
+
+    fn serialize_split(&self, _split: &dyn Split) -> VortexResult<Vec<u8>> {
+        vortex_bail!("LayoutReader splits are not yet serializable");
+    }
+
+    fn deserialize_split(&self, _split: &[u8]) -> VortexResult<SplitRef> {
+        vortex_bail!("LayoutReader splits are not yet serializable");
+    }
+}
+
+struct LayoutReaderScan {
+    dtype: DType,
+    splits: VecDeque<BoxFuture<'static, VortexResult<Option<ArrayRef>>>>,
+}
+
+#[async_trait]
+impl DataSourceScan for LayoutReaderScan {
+    fn dtype(&self) -> &DType {
+        &self.dtype
+    }
+
+    fn remaining_splits_estimate(&self) -> Estimate<usize> {
+        Estimate::Exact(self.splits.len())
+    }
+
+    async fn next_splits(&mut self, max_splits: usize) -> VortexResult<Vec<SplitRef>> {
+        let n = std::cmp::min(max_splits, self.splits.len());
+
+        let mut splits = Vec::with_capacity(n);
+        for _ in 0..n {
+            let fut = self
+                .splits
+                .pop_front()
+                .vortex_expect("Checked length above ensures we have enough splits");
+            splits.push(Box::new(LayoutReaderSplit {
+                dtype: self.dtype.clone(),
+                fut,
+            }) as SplitRef);
+        }
+
+        Ok(splits)
+    }
+}
+
+struct LayoutReaderSplit {
+    dtype: DType,
+    fut: BoxFuture<'static, VortexResult<Option<ArrayRef>>>,
+}
+
+#[async_trait]
+impl Split for LayoutReaderSplit {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn execute(self: Box<Self>) -> VortexResult<SendableArrayStream> {
+        Ok(Box::pin(ArrayStreamAdapter::new(
+            self.dtype,
+            stream::once(self.fut).filter_map(|a| async move { a.transpose() }),
+        )))
+    }
+
+    fn row_count_estimate(&self) -> Estimate<u64> {
+        Estimate::Unknown
+    }
+
+    fn byte_size_estimate(&self) -> Estimate<u64> {
+        Estimate::Unknown
+    }
+}

vortex-scan/src/layout.rs

@@ -6,6 +6,7 @@
 /// We don't actually know if this is right, but it is probably a good estimate.
 const IDEAL_SPLIT_SIZE: u64 = 100_000;
 
+pub mod api;
 pub mod arrow;
 mod filter;
 pub mod row_mask;
@@ -21,6 +22,7 @@ pub use split_by::SplitBy;
 mod scan_builder;
 pub use scan_builder::ScanBuilder;
 
+pub mod layout;
 mod repeated_scan;
 #[cfg(test)]
 mod test;

vortex-scan/src/lib.rs

@@ -56,6 +56,10 @@ pub struct RepeatedScan<A: 'static + Send> {
 }
 
 impl RepeatedScan<ArrayRef> {
+    pub fn dtype(&self) -> &DType {
+        &self.dtype
+    }
+
     pub fn execute_array_iter<B: BlockingRuntime>(
         &self,
         row_range: Option<Range<u64>>,

vortex-scan/src/repeated_scan.rs

@@ -40,6 +40,7 @@ pub mod compressor {
 pub mod dtype {
     pub use vortex_dtype::*;
 }
+
 pub mod error {
     pub use vortex_error::*;
 }