From ca8931277eb1e2fbdd2e6a6178d4d4e8a88e577a Mon Sep 17 00:00:00 2001 From: zhangstar333 Date: Tue, 18 Aug 2026 11:21:18 +0800 Subject: [PATCH] add lance patch --- thirdparty/download-thirdparty.sh | 13 + thirdparty/patches/lance-c-0.1.6-doris.patch | 632 +++++++++++++++++++ 2 files changed, 645 insertions(+) create mode 100644 thirdparty/patches/lance-c-0.1.6-doris.patch diff --git a/thirdparty/download-thirdparty.sh b/thirdparty/download-thirdparty.sh index 7c967d31dd92dd..830da72eacf308 100755 --- a/thirdparty/download-thirdparty.sh +++ b/thirdparty/download-thirdparty.sh @@ -819,6 +819,19 @@ if [[ " ${TP_ARCHIVES[*]} " =~ " PAIMON_CPP " ]]; then echo "Finished patching ${PAIMON_CPP_SOURCE}" fi +# Patch lance-c for fragment-scoped nearest-neighbor search and row-ID-based fetching. +if [[ " ${TP_ARCHIVES[*]} " =~ " LANCE_C " ]]; then + if [[ "${LANCE_C_SOURCE}" == "lance-c-0.1.6" ]]; then + cd "${TP_SOURCE_DIR}/${LANCE_C_SOURCE}" + if [[ ! -f "${PATCHED_MARK}" ]]; then + patch -p1 <"${TP_PATCH_DIR}/lance-c-0.1.6-doris.patch" + touch "${PATCHED_MARK}" + fi + cd - + fi + echo "Finished patching ${LANCE_C_SOURCE}" +fi + if [[ " ${TP_ARCHIVES[*]} " =~ " CCTZ " ]] ; then cd $TP_SOURCE_DIR/$CCTZ_SOURCE if [[ ! -f "$PATCHED_MARK" ]] ; then diff --git a/thirdparty/patches/lance-c-0.1.6-doris.patch b/thirdparty/patches/lance-c-0.1.6-doris.patch new file mode 100644 index 00000000000000..cc49293d9dedf8 --- /dev/null +++ b/thirdparty/patches/lance-c-0.1.6-doris.patch @@ -0,0 +1,632 @@ +diff --git a/include/lance/lance.h b/include/lance/lance.h +index 1de72fa..a0a8b5e 100644 +--- a/include/lance/lance.h ++++ b/include/lance/lance.h +@@ -752,0 +753,25 @@ int32_t lance_dataset_take( ++/** ++ * Take rows by dataset row IDs. ++ * ++ * Row IDs are values from the `_rowid` scanner column, not zero-based row ++ * offsets. They must belong to the same dataset snapshot used for this read. ++ * Missing or deleted row IDs may be omitted from the result. For found rows, ++ * input order and duplicates are preserved. ++ * ++ * @param dataset Open dataset snapshot. ++ * @param row_ids Array of dataset row IDs. May be NULL only when ++ * `num_row_ids` is zero. ++ * @param num_row_ids Length of `row_ids`. ++ * @param columns NULL-terminated column names, or NULL for all. The ++ * system column `_rowid` may be requested explicitly. ++ * @param out Pointer to caller-allocated ArrowArrayStream. ++ * @return 0 on success, -1 on error. ++ */ ++int32_t lance_dataset_take_rows( ++ const LanceDataset* dataset, ++ const uint64_t* row_ids, ++ size_t num_row_ids, ++ const char* const* columns, ++ struct ArrowArrayStream* out ++); ++ +diff --git a/include/lance/lance.hpp b/include/lance/lance.hpp +index 40aa9e3..d54dc5a 100644 +--- a/include/lance/lance.hpp ++++ b/include/lance/lance.hpp +@@ -654,0 +655,24 @@ public: ++ /// Take rows by dataset row IDs. Results exported as ArrowArrayStream. ++ void take_rows(const uint64_t* row_ids, size_t num_row_ids, ++ const std::vector& columns, ++ ArrowArrayStream* out) const { ++ std::vector col_ptrs; ++ for (auto& c : columns) col_ptrs.push_back(c.c_str()); ++ col_ptrs.push_back(nullptr); ++ const char* const* cols_ptr = columns.empty() ? nullptr : col_ptrs.data(); ++ ++ if (lance_dataset_take_rows( ++ handle_.get(), row_ids, num_row_ids, cols_ptr, out) != 0) { ++ check_error(); ++ } ++ } ++ ++ /// Take all columns by dataset row IDs. ++ void take_rows(const uint64_t* row_ids, size_t num_row_ids, ++ ArrowArrayStream* out) const { ++ if (lance_dataset_take_rows( ++ handle_.get(), row_ids, num_row_ids, nullptr, out) != 0) { ++ check_error(); ++ } ++ } ++ +diff --git a/src/dataset.rs b/src/dataset.rs +index 91a735f..364be09 100644 +--- a/src/dataset.rs ++++ b/src/dataset.rs +@@ -44,0 +45,20 @@ impl LanceDataset { ++fn projection_from_columns( ++ dataset: &Dataset, ++ columns: Option<&[String]>, ++) -> Result { ++ match columns { ++ Some(columns) => { ++ let schema = dataset ++ .schema() ++ .project_preserve_system_columns(columns) ++ .map_err(|err| { ++ lance_core::Error::invalid_input(format!("invalid columns {columns:?}: {err}")) ++ })?; ++ Ok(lance::dataset::ProjectionRequest::from_schema(schema)) ++ } ++ None => Ok(lance::dataset::ProjectionRequest::from_schema( ++ dataset.schema().clone(), ++ )), ++ } ++} ++ +@@ -247,4 +267 @@ unsafe fn dataset_take_inner( +- let projection = match &col_names { +- Some(cols) => lance::dataset::ProjectionRequest::from_columns(cols.iter(), snap.schema()), +- None => lance::dataset::ProjectionRequest::from_schema(snap.schema().clone()), +- }; ++ let projection = projection_from_columns(&snap, col_names.as_deref())?; +@@ -263,0 +281,69 @@ unsafe fn dataset_take_inner( ++/// Take rows by dataset row IDs, returning results as an ArrowArrayStream. ++/// ++/// - `row_ids`: array of dataset row IDs, such as values returned in the ++/// `_rowid` scanner column ++/// - `num_row_ids`: length of the row ID array ++/// - `columns`: NULL-terminated column name array, or NULL for all columns ++/// - `out`: pointer to a stack-allocated `ArrowArrayStream` ++/// ++/// `row_ids` may be NULL only when `num_row_ids` is zero. Row IDs must belong ++/// to the same dataset snapshot used for this read. Missing or deleted row IDs ++/// may be omitted from the result by the upstream Lance implementation. ++/// ++/// Returns 0 on success, -1 on error. ++#[unsafe(no_mangle)] ++pub unsafe extern "C" fn lance_dataset_take_rows( ++ dataset: *const LanceDataset, ++ row_ids: *const u64, ++ num_row_ids: usize, ++ columns: *const *const c_char, ++ out: *mut FFI_ArrowArrayStream, ++) -> i32 { ++ ffi_try!( ++ unsafe { dataset_take_rows_inner(dataset, row_ids, num_row_ids, columns, out) }, ++ neg ++ ) ++} ++ ++unsafe fn dataset_take_rows_inner( ++ dataset: *const LanceDataset, ++ row_ids: *const u64, ++ num_row_ids: usize, ++ columns: *const *const c_char, ++ out: *mut FFI_ArrowArrayStream, ++) -> Result { ++ if dataset.is_null() { ++ return Err(lance_core::Error::invalid_input("dataset must not be NULL")); ++ } ++ if out.is_null() { ++ return Err(lance_core::Error::invalid_input("out must not be NULL")); ++ } ++ if num_row_ids > 0 && row_ids.is_null() { ++ return Err(lance_core::Error::invalid_input(format!( ++ "row_ids must not be NULL when num_row_ids = {num_row_ids}" ++ ))); ++ } ++ ++ let ds = unsafe { &*dataset }; ++ let row_id_slice = if num_row_ids == 0 { ++ &[] ++ } else { ++ unsafe { std::slice::from_raw_parts(row_ids, num_row_ids) } ++ }; ++ let col_names = unsafe { helpers::parse_c_string_array(columns)? }; ++ ++ let snap = ds.snapshot(); ++ let projection = projection_from_columns(&snap, col_names.as_deref())?; ++ ++ let batch = block_on(snap.take_rows(row_id_slice, projection))?; ++ ++ // Match lance_dataset_take: export the single RecordBatch as an Arrow stream. ++ let schema = batch.schema(); ++ let reader = arrow::record_batch::RecordBatchIterator::new(vec![Ok(batch)], schema); ++ let ffi_stream = FFI_ArrowArrayStream::new(Box::new(reader)); ++ unsafe { ++ std::ptr::write_unaligned(out, ffi_stream); ++ } ++ Ok(0) ++} ++ +diff --git a/src/scanner.rs b/src/scanner.rs +index c1c2f67..de646be 100644 +--- a/src/scanner.rs ++++ b/src/scanner.rs +@@ -160,0 +161,6 @@ impl LanceScanner { ++ // Lance validates fragment-scoped nearest searches when nearest() is ++ // configured. Such searches are supported when the fragment scan is ++ // the input to a prefilter, so this flag must be set first. ++ if self.prefilter { ++ scanner.prefilter(true); ++ } +@@ -178,3 +183,0 @@ impl LanceScanner { +- if self.prefilter { +- scanner.prefilter(true); +- } +@@ -221,0 +225,5 @@ impl LanceScanner { ++ // nearest() checks the current prefilter setting before accepting a ++ // fragment-scoped search. Enable it before installing the query. ++ if self.prefilter { ++ scanner.prefilter(true); ++ } +@@ -239,3 +246,0 @@ impl LanceScanner { +- if self.prefilter { +- scanner.prefilter(true); +- } +diff --git a/tests/c_api_test.rs b/tests/c_api_test.rs +index 2babda6..9ee5af9 100644 +--- a/tests/c_api_test.rs ++++ b/tests/c_api_test.rs +@@ -9,0 +10 @@ use std::ffi::{CString, c_char}; ++use std::process::Command; +@@ -18 +19 @@ use arrow::record_batch::RecordBatchReader; +-use arrow_array::{Array, Float32Array, Int32Array, RecordBatch, StringArray}; ++use arrow_array::{Array, Float32Array, Int32Array, RecordBatch, StringArray, UInt64Array}; +@@ -393,0 +395,94 @@ fn test_dataset_take() { ++#[test] ++fn test_dataset_take_rows_empty_and_null_validation() { ++ let (_tmp, uri) = create_test_dataset(); ++ let c_uri = c_str(&uri); ++ let ds = unsafe { lance_dataset_open(c_uri.as_ptr(), ptr::null(), 0) }; ++ assert!(!ds.is_null()); ++ ++ let mut empty_stream = FFI_ArrowArrayStream::empty(); ++ assert_eq!( ++ unsafe { lance_dataset_take_rows(ds, ptr::null(), 0, ptr::null(), &mut empty_stream) }, ++ 0 ++ ); ++ let reader = unsafe { ArrowArrayStreamReader::from_raw(&mut empty_stream) }.unwrap(); ++ assert_eq!( ++ reader.map(|batch| batch.unwrap().num_rows()).sum::(), ++ 0 ++ ); ++ ++ let mut invalid_stream = FFI_ArrowArrayStream::empty(); ++ assert_eq!( ++ unsafe { lance_dataset_take_rows(ds, ptr::null(), 1, ptr::null(), &mut invalid_stream) }, ++ -1 ++ ); ++ let message = unsafe { std::ffi::CStr::from_ptr(lance_last_error_message()) }.to_string_lossy(); ++ assert!( ++ message.contains("row_ids must not be NULL when num_row_ids = 1"), ++ "unexpected error: {message}" ++ ); ++ ++ let row_id = 0_u64; ++ assert_eq!( ++ unsafe { ++ lance_dataset_take_rows(ptr::null(), &row_id, 1, ptr::null(), &mut invalid_stream) ++ }, ++ -1 ++ ); ++ assert_eq!( ++ unsafe { lance_dataset_take_rows(ds, &row_id, 1, ptr::null(), ptr::null_mut()) }, ++ -1 ++ ); ++ ++ unsafe { lance_dataset_close(ds) }; ++} ++ ++#[test] ++fn test_dataset_take_rows_invalid_column() { ++ const CHILD_ENV: &str = "LANCE_C_TEST_TAKE_ROWS_INVALID_COLUMN_CHILD"; ++ ++ if std::env::var_os(CHILD_ENV).is_some() { ++ let (_tmp, uri) = create_test_dataset(); ++ let c_uri = c_str(&uri); ++ let ds = unsafe { lance_dataset_open(c_uri.as_ptr(), ptr::null(), 0) }; ++ assert!(!ds.is_null()); ++ ++ let row_id = 0_u64; ++ let invalid_column = c_str("unknown_column"); ++ let columns = [invalid_column.as_ptr(), ptr::null()]; ++ let mut stream = FFI_ArrowArrayStream::empty(); ++ assert_eq!( ++ unsafe { lance_dataset_take_rows(ds, &row_id, 1, columns.as_ptr(), &mut stream) }, ++ -1 ++ ); ++ assert_eq!(lance_last_error_code(), LanceErrorCode::InvalidArgument); ++ ++ let message_ptr = lance_last_error_message(); ++ assert!(!message_ptr.is_null()); ++ let message = unsafe { std::ffi::CStr::from_ptr(message_ptr) } ++ .to_string_lossy() ++ .into_owned(); ++ unsafe { lance_free_string(message_ptr) }; ++ assert!( ++ message.contains("unknown_column"), ++ "unexpected error: {message}" ++ ); ++ ++ unsafe { lance_dataset_close(ds) }; ++ return; ++ } ++ ++ let output = Command::new(std::env::current_exe().unwrap()) ++ .arg("--exact") ++ .arg("test_dataset_take_rows_invalid_column") ++ .arg("--nocapture") ++ .env(CHILD_ENV, "1") ++ .output() ++ .unwrap(); ++ assert!( ++ output.status.success(), ++ "invalid-column subprocess failed\nstdout:\n{}\nstderr:\n{}", ++ String::from_utf8_lossy(&output.stdout), ++ String::from_utf8_lossy(&output.stderr) ++ ); ++} ++ +@@ -2497,0 +2593,73 @@ fn create_vector_dataset(num_rows: i32, dim: i32) -> (tempfile::TempDir, String) ++/// Create two vector fragments with deterministic vectors. Every component of ++/// row `id` is `id as f32`, making nearest-neighbor expectations unambiguous. ++fn create_multi_fragment_vector_dataset( ++ rows_per_fragment: i32, ++ dim: i32, ++ enable_stable_row_ids: bool, ++) -> (tempfile::TempDir, String) { ++ use arrow_array::builder::{FixedSizeListBuilder, Float32Builder}; ++ ++ let tmp = tempfile::tempdir().unwrap(); ++ let uri = tmp ++ .path() ++ .join("multi_fragment_vec_ds") ++ .to_str() ++ .unwrap() ++ .to_string(); ++ let schema = Arc::new(Schema::new(vec![ ++ Field::new("id", DataType::Int32, false), ++ Field::new( ++ "embedding", ++ DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), ++ false, ++ ), ++ ])); ++ ++ let make_batch = |first_id: i32| { ++ let ids: Vec = (first_id..first_id + rows_per_fragment).collect(); ++ let mut embeddings = FixedSizeListBuilder::new(Float32Builder::new(), dim); ++ for id in &ids { ++ for _ in 0..dim { ++ embeddings.values().append_value(*id as f32); ++ } ++ embeddings.append(true); ++ } ++ RecordBatch::try_new( ++ schema.clone(), ++ vec![ ++ Arc::new(Int32Array::from(ids)), ++ Arc::new(embeddings.finish()), ++ ], ++ ) ++ .unwrap() ++ }; ++ ++ let first = make_batch(0); ++ let second = make_batch(rows_per_fragment); ++ lance_c::runtime::block_on(async { ++ Dataset::write( ++ arrow::record_batch::RecordBatchIterator::new(vec![Ok(first)], schema.clone()), ++ &uri, ++ Some(lance::dataset::WriteParams { ++ enable_stable_row_ids, ++ ..Default::default() ++ }), ++ ) ++ .await ++ .unwrap(); ++ Dataset::write( ++ arrow::record_batch::RecordBatchIterator::new(vec![Ok(second)], schema), ++ &uri, ++ Some(lance::dataset::WriteParams { ++ mode: lance::dataset::WriteMode::Append, ++ enable_stable_row_ids, ++ ..Default::default() ++ }), ++ ) ++ .await ++ .unwrap(); ++ }); ++ ++ (tmp, uri) ++} ++ +@@ -2723,0 +2892,99 @@ fn test_scanner_nearest_brute_force() { ++fn assert_dataset_take_rows_from_multi_fragment_ann_result(enable_stable_row_ids: bool) { ++ let (_tmp, uri) = create_multi_fragment_vector_dataset(32, 8, enable_stable_row_ids); ++ let uri_c = c_str(&uri); ++ let ds = unsafe { lance_dataset_open(uri_c.as_ptr(), ptr::null(), 0) }; ++ assert!(!ds.is_null()); ++ ++ // A non-NULL array whose first element is NULL is an explicit empty ++ // projection. The ANN result should therefore contain only _distance and ++ // the explicitly requested _rowid system column. ++ let no_columns: [*const c_char; 1] = [ptr::null()]; ++ let scanner = unsafe { lance_scanner_new(ds, no_columns.as_ptr(), ptr::null()) }; ++ assert!(!scanner.is_null()); ++ assert_eq!(unsafe { lance_scanner_with_row_id(scanner, true) }, 0); ++ ++ let column = c_str("embedding"); ++ let query = [40.0_f32; 8]; ++ assert_eq!( ++ unsafe { ++ lance_scanner_nearest( ++ scanner, ++ column.as_ptr(), ++ query.as_ptr().cast(), ++ query.len(), ++ LanceDataType::Float32 as i32, ++ 1, ++ ) ++ }, ++ 0 ++ ); ++ assert_eq!(unsafe { lance_scanner_set_use_index(scanner, false) }, 0); ++ ++ let mut ann_stream = FFI_ArrowArrayStream::empty(); ++ assert_eq!( ++ unsafe { lance_scanner_to_arrow_stream(scanner, &mut ann_stream) }, ++ 0 ++ ); ++ let reader = unsafe { ArrowArrayStreamReader::from_raw(&mut ann_stream) }.unwrap(); ++ let ann_batches = reader.map(|batch| batch.unwrap()).collect::>(); ++ assert_eq!( ++ ann_batches.iter().map(RecordBatch::num_rows).sum::(), ++ 1 ++ ); ++ assert_eq!(ann_batches[0].num_columns(), 2); ++ ++ let distance = ann_batches[0] ++ .column_by_name("_distance") ++ .unwrap() ++ .as_any() ++ .downcast_ref::() ++ .unwrap() ++ .value(0); ++ assert_eq!(distance, 0.0); ++ let row_id = ann_batches[0] ++ .column_by_name("_rowid") ++ .unwrap() ++ .as_any() ++ .downcast_ref::() ++ .unwrap() ++ .value(0); ++ if !enable_stable_row_ids { ++ assert_ne!( ++ row_id >> 32, ++ 0, ++ "expected an address-style row ID from the second fragment" ++ ); ++ } ++ ++ let id_column = c_str("id"); ++ let columns = [id_column.as_ptr(), ptr::null()]; ++ let mut take_stream = FFI_ArrowArrayStream::empty(); ++ assert_eq!( ++ unsafe { lance_dataset_take_rows(ds, &row_id, 1, columns.as_ptr(), &mut take_stream) }, ++ 0 ++ ); ++ let reader = unsafe { ArrowArrayStreamReader::from_raw(&mut take_stream) }.unwrap(); ++ let batches = reader.map(|batch| batch.unwrap()).collect::>(); ++ assert_eq!(batches.len(), 1); ++ let ids = batches[0] ++ .column_by_name("id") ++ .unwrap() ++ .as_any() ++ .downcast_ref::() ++ .unwrap(); ++ assert_eq!(ids.values(), &[40]); ++ ++ unsafe { lance_scanner_close(scanner) }; ++ unsafe { lance_dataset_close(ds) }; ++} ++ ++#[test] ++fn test_dataset_take_rows_from_multi_fragment_ann_result() { ++ assert_dataset_take_rows_from_multi_fragment_ann_result(false); ++} ++ ++#[test] ++fn test_dataset_take_rows_from_multi_fragment_ann_result_with_stable_row_ids() { ++ assert_dataset_take_rows_from_multi_fragment_ann_result(true); ++} ++ +@@ -2855,0 +3123,132 @@ fn test_scanner_nearest_filter_postfilter() { ++#[test] ++fn test_scanner_nearest_prefilter_with_fragment_ids_next() { ++ let (_tmp, uri) = create_multi_fragment_vector_dataset(32, 8, false); ++ let uri_c = c_str(&uri); ++ let ds = unsafe { lance_dataset_open(uri_c.as_ptr(), ptr::null(), 0) }; ++ assert!(!ds.is_null()); ++ ++ let mut fragment_ids = vec![0; unsafe { lance_dataset_fragment_count(ds) } as usize]; ++ assert_eq!(fragment_ids.len(), 2); ++ assert_eq!( ++ unsafe { lance_dataset_fragment_ids(ds, fragment_ids.as_mut_ptr()) }, ++ 0 ++ ); ++ ++ let filter = c_str("id >= 40"); ++ let scanner = unsafe { lance_scanner_new(ds, ptr::null(), filter.as_ptr()) }; ++ assert_eq!( ++ unsafe { lance_scanner_set_fragment_ids(scanner, fragment_ids[1..].as_ptr(), 1) }, ++ 0 ++ ); ++ ++ // Match the Doris call order: nearest is configured before prefilter. ++ let column = c_str("embedding"); ++ let query = [40.0_f32; 8]; ++ assert_eq!( ++ unsafe { ++ lance_scanner_nearest( ++ scanner, ++ column.as_ptr(), ++ query.as_ptr().cast(), ++ query.len(), ++ LanceDataType::Float32 as i32, ++ 5, ++ ) ++ }, ++ 0 ++ ); ++ assert_eq!(unsafe { lance_scanner_set_prefilter(scanner, true) }, 0); ++ assert_eq!(unsafe { lance_scanner_set_use_index(scanner, false) }, 0); ++ ++ let batches = scan_all_rows_from_scanner(scanner); ++ let mut ids = batches ++ .iter() ++ .flat_map(|batch| { ++ batch ++ .column_by_name("id") ++ .unwrap() ++ .as_any() ++ .downcast_ref::() ++ .unwrap() ++ .values() ++ .iter() ++ .copied() ++ .collect::>() ++ }) ++ .collect::>(); ++ ids.sort_unstable(); ++ assert_eq!(ids, vec![40, 41, 42, 43, 44]); ++ ++ unsafe { lance_scanner_close(scanner) }; ++ unsafe { lance_dataset_close(ds) }; ++} ++ ++#[test] ++fn test_scanner_nearest_prefilter_with_fragment_ids_arrow_stream() { ++ let (_tmp, uri) = create_multi_fragment_vector_dataset(32, 8, false); ++ let uri_c = c_str(&uri); ++ let ds = unsafe { lance_dataset_open(uri_c.as_ptr(), ptr::null(), 0) }; ++ assert!(!ds.is_null()); ++ ++ let mut fragment_ids = vec![0; unsafe { lance_dataset_fragment_count(ds) } as usize]; ++ assert_eq!(fragment_ids.len(), 2); ++ assert_eq!( ++ unsafe { lance_dataset_fragment_ids(ds, fragment_ids.as_mut_ptr()) }, ++ 0 ++ ); ++ ++ let filter = c_str("id >= 60"); ++ let scanner = unsafe { lance_scanner_new(ds, ptr::null(), filter.as_ptr()) }; ++ assert_eq!( ++ unsafe { lance_scanner_set_fragment_ids(scanner, fragment_ids[1..].as_ptr(), 1) }, ++ 0 ++ ); ++ ++ let column = c_str("embedding"); ++ let query = [60.0_f32; 8]; ++ assert_eq!( ++ unsafe { ++ lance_scanner_nearest( ++ scanner, ++ column.as_ptr(), ++ query.as_ptr().cast(), ++ query.len(), ++ LanceDataType::Float32 as i32, ++ 10, ++ ) ++ }, ++ 0 ++ ); ++ assert_eq!(unsafe { lance_scanner_set_prefilter(scanner, true) }, 0); ++ assert_eq!(unsafe { lance_scanner_set_use_index(scanner, false) }, 0); ++ ++ let mut stream = FFI_ArrowArrayStream::empty(); ++ assert_eq!( ++ unsafe { lance_scanner_to_arrow_stream(scanner, &mut stream) }, ++ 0, ++ "{}", ++ unsafe { std::ffi::CStr::from_ptr(lance_last_error_message()) }.to_string_lossy() ++ ); ++ let reader = unsafe { ArrowArrayStreamReader::from_raw(&mut stream) }.unwrap(); ++ let mut ids = reader ++ .flat_map(|batch| { ++ let batch = batch.unwrap(); ++ batch ++ .column_by_name("id") ++ .unwrap() ++ .as_any() ++ .downcast_ref::() ++ .unwrap() ++ .values() ++ .iter() ++ .copied() ++ .collect::>() ++ }) ++ .collect::>(); ++ ids.sort_unstable(); ++ assert_eq!(ids, vec![60, 61, 62, 63]); ++ ++ unsafe { lance_scanner_close(scanner) }; ++ unsafe { lance_dataset_close(ds) }; ++} ++ +diff --git a/tests/cpp/test_cpp_api.cpp b/tests/cpp/test_cpp_api.cpp +index 43491f2..444fbdc 100644 +--- a/tests/cpp/test_cpp_api.cpp ++++ b/tests/cpp/test_cpp_api.cpp +@@ -124,0 +125,29 @@ static void test_dataset_take(const std::string& uri) { ++static void test_dataset_take_rows(const std::string& uri) { ++ TEST(test_dataset_take_rows); ++ ++ auto ds = lance::Dataset::open(uri); ++ ++ // The smoke fixture has one fragment, so its first row IDs are 0, 1, 2. ++ uint64_t row_ids[] = {0, 1, 2}; ++ ArrowArrayStream stream; ++ memset(&stream, 0, sizeof(stream)); ++ ds.take_rows(row_ids, 3, &stream); ++ ++ uint64_t total = 0; ++ while (true) { ++ ArrowArray arr; ++ memset(&arr, 0, sizeof(arr)); ++ int rc = stream.get_next(&stream, &arr); ++ assert(rc == 0); ++ if (!arr.release) break; ++ total += (uint64_t)arr.length; ++ arr.release(&arr); ++ } ++ ++ assert(total == 3); ++ printf("rows=%llu... ", (unsigned long long)total); ++ ++ if (stream.release) stream.release(&stream); ++ PASS(); ++} ++ +@@ -695,0 +725 @@ int main(int argc, char** argv) { ++ test_dataset_take_rows(uri);