Rust bindings for CDF reads #612

PatrickJin-db · 2024-12-11T08:02:44Z

To be released as delta-kernel-rust-sharing-wrapper 0.2.0

linzhou-db

I'll leave the main review to the kernel rust team

linzhou-db · 2024-12-11T18:38:13Z

python/delta-kernel-rust-sharing-wrapper/.gitignore

@@ -0,0 +1 @@
+Cargo.lock


What's this for?

When you run cargo commands or maturin develop, a Cargo.lock file may be generated. Adding it to the gitignore makes sure you don't accidentally add it to git.

for further context, generally for libraries (like this) you don't want to commit whereas binaries you normally do

linzhou-db · 2024-12-11T18:38:38Z

python/delta-kernel-rust-sharing-wrapper/src/lib.rs

+    })?))
+}
+
+// TODO(patrick.jin): Change return type to RecordBatchIterator<impl Iterator<...>>


should we finish this TODO before merging?

I tried to but it seems non-trivial. But I will try to address it before the final release.

More specifically, changing the Vec to an Iterator requires the return lifetime of Scan::execute and TableChangesScan::execute to be '_, and that requires the return type of those functions to be ...<Arc<...>> insteadof ...<Box<...>>. And upon making that change, I get a cannot be converted to a python object error. Will see if @zachschuermann or @OussamaSaoudi-db knows more about this error.

Posted a fix this morning, so this is now fixed :)

OussamaSaoudi-db · 2024-12-11T20:08:22Z

python/delta-kernel-rust-sharing-wrapper/src/lib.rs

+fn try_create_record_batch_iter(
+    results: impl Iterator<Item = DeltaResult<ScanResult>>,
+    result_schema: ArrowSchemaRef,
+) -> RecordBatchIterator<Vec<Result<RecordBatch, ArrowError>>> {
+    let record_batches: Vec<_> = results
+        .map(|res| {
+            let scan_res = res.and_then(|res| Ok((res.full_mask(), res.raw_data?)));
+            let (mask, data) =
+                scan_res.map_err(|e| ArrowError::from_external_error(Box::new(e)))?;
+            let record_batch: RecordBatch = data
+                .into_any()
+                .downcast::<ArrowEngineData>()
+                .map_err(|_| ArrowError::CastError("Couldn't cast to ArrowEngineData".to_string()))?
+                .into();
+            if let Some(mask) = mask {
+                let filtered_batch = filter_record_batch(&record_batch, &mask.into())?;
+                Ok(filtered_batch)
+            } else {
+                Ok(record_batch)
+            }
+        })
+        .collect();
+    RecordBatchIterator::new(record_batches, result_schema)
+}
+


Suggested change

fn try_create_record_batch_iter(

results: impl Iterator<Item = DeltaResult<ScanResult>>,

result_schema: ArrowSchemaRef,

) -> RecordBatchIterator<Vec<Result<RecordBatch, ArrowError>>> {

let record_batches: Vec<_> = results

.map(|res| {

let scan_res = res.and_then(|res| Ok((res.full_mask(), res.raw_data?)));

let (mask, data) =

scan_res.map_err(|e| ArrowError::from_external_error(Box::new(e)))?;

let record_batch: RecordBatch = data

.into_any()

.downcast::<ArrowEngineData>()

.map_err(|_| ArrowError::CastError("Couldn't cast to ArrowEngineData".to_string()))?

.into();

if let Some(mask) = mask {

let filtered_batch = filter_record_batch(&record_batch, &mask.into())?;

Ok(filtered_batch)

} else {

Ok(record_batch)

}

})

.collect();

RecordBatchIterator::new(record_batches, result_schema)

}

fn try_create_record_batch_iter(

results: impl Iterator<Item = DeltaResult<ScanResult>>,

result_schema: ArrowSchemaRef,

) -> RecordBatchIterator<impl Iterator<Item = Result<RecordBatch, ArrowError>>> {

let record_batches = results.map(|res| {

let scan_res = res.and_then(|res| Ok((res.full_mask(), res.raw_data?)));

let (mask, data) = scan_res.map_err(|e| ArrowError::from_external_error(Box::new(e)))?;

let record_batch: RecordBatch = data

.into_any()

.downcast::<ArrowEngineData>()

.map_err(|_| ArrowError::CastError("Couldn't cast to ArrowEngineData".to_string()))?

.into();

match mask {

Some(mask) => filter_record_batch(&record_batch, &mask.into()),

None => Ok(record_batch),

}

});

RecordBatchIterator::new(record_batches, result_schema)

}

@OussamaSaoudi-db since this is a pretty common pattern I wonder if somewhere in default enging/arrow implementation we should basically provide this

Yeah, let's create an issue for it. we have it in like 3 tests and our examples, and here, so clearly it would be useful to have a canonical implementation

delta-io/delta-kernel-rs#592

OussamaSaoudi-db · 2024-12-11T20:20:43Z

python/delta-kernel-rust-sharing-wrapper/src/lib.rs

+    }
+
+    fn build(&mut self) -> DeltaPyResult<TableChangesScan> {
+        let scan = self.0.take().unwrap().build()?;


This will panic on unwrap if build is called twice. Perhaps something like

let scan = self .0 .take() .ok_or_else(|| { delta_kernel::Error::generic("Can only call build once on TableChangesScanBuilder") })? .build()?;

And you can put in a suitable error message.

Same issue with ScanBuilder::build

zachschuermann

flushing comments

zachschuermann · 2024-12-11T20:24:09Z

python/delta-kernel-rust-sharing-wrapper/.gitignore

@@ -0,0 +1 @@
+Cargo.lock


for further context, generally for libraries (like this) you don't want to commit whereas binaries you normally do

zachschuermann · 2024-12-11T20:25:13Z

python/delta-kernel-rust-sharing-wrapper/Cargo.toml

@@ -11,7 +11,7 @@ crate-type = ["cdylib"]

 [dependencies]
 arrow = { version = "53.3.0", features = ["pyarrow"] }
-delta_kernel = {version = "0.5", features = ["cloud", "default", "default-engine"]}
+delta_kernel = { git = "https://github.com/delta-io/delta-kernel-rs.git", rev = "be1453f", features = ["cloud", "default", "default-engine"]}


default features are empty

Suggested change

delta_kernel = { git = "https://github.com/delta-io/delta-kernel-rs.git", rev = "be1453f", features = ["cloud", "default", "default-engine"]}

delta_kernel = { git = "https://github.com/delta-io/delta-kernel-rs.git", rev = "be1453f", features = ["cloud", "default-engine"]}

zachschuermann · 2024-12-11T20:32:00Z

python/delta-kernel-rust-sharing-wrapper/src/lib.rs

+fn try_create_record_batch_iter(
+    results: impl Iterator<Item = DeltaResult<ScanResult>>,
+    result_schema: ArrowSchemaRef,
+) -> RecordBatchIterator<Vec<Result<RecordBatch, ArrowError>>> {
+    let record_batches: Vec<_> = results
+        .map(|res| {
+            let scan_res = res.and_then(|res| Ok((res.full_mask(), res.raw_data?)));
+            let (mask, data) =
+                scan_res.map_err(|e| ArrowError::from_external_error(Box::new(e)))?;
+            let record_batch: RecordBatch = data
+                .into_any()
+                .downcast::<ArrowEngineData>()
+                .map_err(|_| ArrowError::CastError("Couldn't cast to ArrowEngineData".to_string()))?
+                .into();
+            if let Some(mask) = mask {
+                let filtered_batch = filter_record_batch(&record_batch, &mask.into())?;
+                Ok(filtered_batch)
+            } else {
+                Ok(record_batch)
+            }
+        })
+        .collect();
+    RecordBatchIterator::new(record_batches, result_schema)
+}
+


@OussamaSaoudi-db since this is a pretty common pattern I wonder if somewhere in default enging/arrow implementation we should basically provide this

OussamaSaoudi-db

I've just got nits on how you may want to refactor/change the code. At a high level, this LGTM

OussamaSaoudi-db · 2024-12-12T21:55:53Z

python/delta-kernel-rust-sharing-wrapper/src/lib.rs

+            .0
+            .take()
+            .ok_or_else(|| {
+                delta_kernel::Error::generic("Can only call build() once on ScanBuilder")


nit: to reduce noise, I'd recommend using type aliases.

use delta_kernel::Error as KernelError; struct PyKernelError(KernelError);

then you would have both

impl From<PyKernelError> for PyErr { fn from(error: PyKernelError) -> Self { PyValueError::new_err(format!("Kernel error: {}", error.0)) } } impl From<KernelError> for PyKernelError { fn from(delta_kernel_error: delta_kernel::Error) -> Self { Self(delta_kernel_error) } }

similar for other delta_kernel::.*

Done for some (but not all) delta_kernel::*. I think refactoring to adopt this pattern everywhere can be a followup

OussamaSaoudi-db · 2024-12-12T22:09:39Z

python/delta-kernel-rust-sharing-wrapper/src/lib.rs

+struct TableChangesScanBuilder(Option<delta_kernel::table_changes::scan::TableChangesScanBuilder>);
+
+#[pymethods]
+impl TableChangesScanBuilder {
+    #[new]
+    #[pyo3(signature = (table, engine_interface, start_version, end_version=None))]
+    fn new(
+        table: &Table,
+        engine_interface: &PythonInterface,
+        start_version: u64,
+        end_version: Option<u64>,
+    ) -> DeltaPyResult<TableChangesScanBuilder> {
+        let table_changes =
+            table
+                .0
+                .table_changes(engine_interface.0.as_ref(), start_version, end_version)?;
+        Ok(TableChangesScanBuilder(Some(
+            table_changes.into_scan_builder(),
+        )))
+    }
+
+    fn build(&mut self) -> DeltaPyResult<TableChangesScan> {
+        let scan = self
+            .0
+            .take()
+            .ok_or_else(|| {
+                delta_kernel::Error::generic(
+                    "Can only call build() once on TableChangesScanBuilder",
+                )
+            })?
+            .build()?;
+        Ok(TableChangesScan(scan))
+    }
+}


Currently we fail on any subsequent call to build. If we want to avoid the potential error, you can do something like this:

Suggested change

struct TableChangesScanBuilder(Option<delta_kernel::table_changes::scan::TableChangesScanBuilder>);

#[pymethods]

impl TableChangesScanBuilder {

#[new]

#[pyo3(signature = (table, engine_interface, start_version, end_version=None))]

fn new(

table: &Table,

engine_interface: &PythonInterface,

start_version: u64,

end_version: Option<u64>,

) -> DeltaPyResult<TableChangesScanBuilder> {

let table_changes =

table

.0

.table_changes(engine_interface.0.as_ref(), start_version, end_version)?;

Ok(TableChangesScanBuilder(Some(

table_changes.into_scan_builder(),

)))

}

fn build(&mut self) -> DeltaPyResult<TableChangesScan> {

let scan = self

.0

.take()

.ok_or_else(|| {

delta_kernel::Error::generic(

"Can only call build() once on TableChangesScanBuilder",

)

})?

.build()?;

Ok(TableChangesScan(scan))

}

}

#[pyclass]

struct TableChangesScanBuilder {

table_changes: Arc<TableChanges>,

}

#[pymethods]

impl TableChangesScanBuilder {

#[new]

#[pyo3(signature = (table, engine_interface, start_version, end_version=None))]

fn new(

table: &Table,

engine_interface: &PythonInterface,

start_version: u64,

end_version: Option<u64>,

) -> DeltaPyResult<TableChangesScanBuilder> {

let table_changes = table

.0

.table_changes(engine_interface.0.as_ref(), start_version, end_version)?

.into();

Ok(TableChangesScanBuilder { table_changes })

}

fn build(&mut self) -> DeltaPyResult<TableChangesScan> {

// Note: cheap Arc clone

let scan_builder = self.table_changes.clone().scan_builder();

let scan = scan_builder.build()?;

Ok(TableChangesScan(scan))

}

}

You can do the same thing in Snapsot/ScanBulider. @zachschuermann would like to see what you think.

I think if we change this behavior, we should change it for both TableChangesScanBuilder and ScanBuilder to be consistent.

python/delta-kernel-rust-sharing-wrapper/src/lib.rs

nicklan · 2024-12-15T00:25:48Z

python/delta-kernel-rust-sharing-wrapper/src/lib.rs

        Ok(PyArrowType(Box::new(record_batch_iter)))
    }
 }

 #[pyclass]
-struct PythonInterface(Box<dyn Engine + Send>);
+struct PythonInterface(Arc<dyn Engine + Send>);


what was the motivation to change from box to arc?

It is related to this kernel change. Now we have to pass an Arc<dyn Engine> into Scan::execute/TableChangesScan::execute.

Co-authored-by: Oussama Saoudi <[email protected]>

PatrickJin-db requested a review from zachschuermann December 11, 2024 08:02

linzhou-db requested a review from OussamaSaoudi-db December 11, 2024 18:38

linzhou-db reviewed Dec 11, 2024

View reviewed changes

PatrickJin-db force-pushed the PatrickJin-db/cdf_rust_bindings branch from 7dfe565 to 5ce85e1 Compare December 11, 2024 19:37

OussamaSaoudi-db reviewed Dec 11, 2024

View reviewed changes

PatrickJin-db force-pushed the PatrickJin-db/cdf_rust_bindings branch 4 times, most recently from 8f0d513 to aaf3e43 Compare December 11, 2024 20:34

zachschuermann requested a review from nicklan December 11, 2024 20:37

zachschuermann reviewed Dec 11, 2024

View reviewed changes

PatrickJin-db force-pushed the PatrickJin-db/cdf_rust_bindings branch 5 times, most recently from df5ab92 to 6351118 Compare December 12, 2024 02:39

PatrickJin-db requested review from zachschuermann and OussamaSaoudi-db December 12, 2024 20:30

OussamaSaoudi-db approved these changes Dec 12, 2024

View reviewed changes

PatrickJin-db force-pushed the PatrickJin-db/cdf_rust_bindings branch 5 times, most recently from 979c558 to 757aa83 Compare December 13, 2024 09:41

nicklan reviewed Dec 15, 2024

View reviewed changes

PatrickJin-db force-pushed the PatrickJin-db/cdf_rust_bindings branch 2 times, most recently from ac1659c to 6449594 Compare December 18, 2024 01:45

rust bindings for CDF

04c7a31

Co-authored-by: Oussama Saoudi <[email protected]>

PatrickJin-db force-pushed the PatrickJin-db/cdf_rust_bindings branch from 6449594 to 04c7a31 Compare December 18, 2024 02:20

PatrickJin-db merged commit 9c33111 into delta-io:main Dec 18, 2024
14 checks passed

PatrickJin-db mentioned this pull request Dec 20, 2024

update delta-kernel-rust-sharing-wrapper to 0.2.1 and fix path in workflows #625

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Rust bindings for CDF reads #612

Rust bindings for CDF reads #612

PatrickJin-db commented Dec 11, 2024 •

edited

Loading

linzhou-db left a comment

linzhou-db Dec 11, 2024

PatrickJin-db Dec 11, 2024

zachschuermann Dec 11, 2024

linzhou-db Dec 11, 2024

PatrickJin-db Dec 11, 2024

OussamaSaoudi-db Dec 11, 2024

OussamaSaoudi-db Dec 11, 2024 •

edited

Loading

zachschuermann Dec 11, 2024

nicklan Dec 11, 2024

zachschuermann Dec 12, 2024

OussamaSaoudi-db Dec 11, 2024

OussamaSaoudi-db Dec 11, 2024

zachschuermann left a comment

zachschuermann Dec 11, 2024

zachschuermann Dec 11, 2024

zachschuermann Dec 11, 2024

OussamaSaoudi-db left a comment

OussamaSaoudi-db Dec 12, 2024

OussamaSaoudi-db Dec 12, 2024

PatrickJin-db Dec 13, 2024

OussamaSaoudi-db Dec 12, 2024

OussamaSaoudi-db Dec 12, 2024

PatrickJin-db Dec 13, 2024

nicklan Dec 15, 2024

PatrickJin-db Dec 16, 2024

	delta_kernel = { git = "https://github.com/delta-io/delta-kernel-rs.git", rev = "be1453f", features = ["cloud", "default", "default-engine"]}
	delta_kernel = { git = "https://github.com/delta-io/delta-kernel-rs.git", rev = "be1453f", features = ["cloud", "default-engine"]}

		@@ -0,0 +1 @@
		Cargo.lock

		@@ -0,0 +1 @@
		Cargo.lock

Rust bindings for CDF reads #612

Rust bindings for CDF reads #612

Conversation

PatrickJin-db commented Dec 11, 2024 • edited Loading

linzhou-db left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

OussamaSaoudi-db Dec 11, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

zachschuermann left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

OussamaSaoudi-db left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

PatrickJin-db commented Dec 11, 2024 •

edited

Loading

OussamaSaoudi-db Dec 11, 2024 •

edited

Loading