diff --git a/java/src/test/java/org/lance/ScannerTest.java b/java/src/test/java/org/lance/ScannerTest.java index 00434034b64..0b07026bc68 100644 --- a/java/src/test/java/org/lance/ScannerTest.java +++ b/java/src/test/java/org/lance/ScannerTest.java @@ -422,25 +422,32 @@ void testDatasetScannerBatchReadahead(@TempDir Path tempDir) throws Exception { TestUtils.SimpleTestDataset testDataset = new TestUtils.SimpleTestDataset(allocator, datasetPath); testDataset.createEmptyDataset().close(); - int totalRows = 1000; - int batchSize = 100; - int batchReadahead = 5; - try (Dataset dataset = testDataset.write(1, totalRows)) { + + int totalRows = 2000; + int maxRowsPerFile = 100; // ~20 fragments + List fragments = testDataset.createNewFragment(totalRows, maxRowsPerFile); + assertTrue(fragments.size() > 1, "expected multiple fragments, got " + fragments.size()); + + FragmentOperation.Append append = new FragmentOperation.Append(fragments); + try (Dataset dataset = Dataset.commit(allocator, datasetPath, append, Optional.of(1L))) { + int batchReadahead = 2; // far below the default (num compute CPUs) try (LanceScanner scanner = dataset.newScan( - new ScanOptions.Builder() - .batchSize(batchSize) - .batchReadahead(batchReadahead) - .build())) { - // This test is more about ensuring that the batchReadahead parameter is accepted - // and doesn't cause errors. The actual effect of batchReadahead might not be - // directly observable in this test. + new ScanOptions.Builder().batchSize(50).batchReadahead(batchReadahead).build())) { try (ArrowReader reader = scanner.scanBatches()) { int rowCount = 0; + long idSum = 0; while (reader.loadNextBatch()) { - rowCount += reader.getVectorSchemaRoot().getRowCount(); + VectorSchemaRoot root = reader.getVectorSchemaRoot(); + IntVector ids = (IntVector) root.getVector("id"); + for (int i = 0; i < root.getRowCount(); i++) { + idSum += ids.get(i); + } + rowCount += root.getRowCount(); } assertEquals(totalRows, rowCount); + // ids are the contiguous range [0, totalRows) + assertEquals((long) totalRows * (totalRows - 1) / 2, idSum); } } } diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 7c7a082ad7e..d80ce3214e4 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -91,7 +91,9 @@ use crate::index::scalar_logical::scalar_index_fragment_bitmap; use crate::index::vector::utils::{ default_distance_type_for, get_vector_dim, get_vector_type, validate_distance_type_for, }; -use crate::io::exec::filtered_read::{FilteredReadExec, FilteredReadOptions}; +use crate::io::exec::filtered_read::{ + FilteredReadExec, FilteredReadOptions, FilteredReadThreadingMode, +}; use crate::io::exec::fts::{ BoostQueryExec, FlatMatchFilterExec, FlatMatchQueryExec, MatchQueryExec, PhraseQueryExec, }; @@ -1374,8 +1376,10 @@ impl Scanner { self } - /// Set the prefetch size. - /// Ignored in v2 and newer format + /// Set the number of batches to decode concurrently. + /// + /// This bounds the decode fan-out of the scan: at most this many batch-decode + /// tasks run in flight at once. Defaults to `get_num_compute_intensive_cpus()`. pub fn batch_readahead(&mut self, nbatches: usize) -> &mut Self { self.batch_readahead = nbatches; self @@ -2877,6 +2881,11 @@ impl Scanner { read_options = read_options.with_batch_size(batch_size as u32); } + // Bound the decode fan-out by `batch_readahead`. + read_options = read_options.with_threading_mode( + FilteredReadThreadingMode::OnePartitionMultipleThreads(self.batch_readahead), + ); + if let Some(file_reader_options) = self.resolved_file_reader_options() { read_options = read_options.with_file_reader_options(file_reader_options); } @@ -12374,6 +12383,36 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") assert_eq!(filtered.options().io_buffer_size_bytes, Some(7777)); } + #[tokio::test] + async fn test_batch_readahead_bounds_decode_concurrency() { + let data = lance_datagen::gen_batch() + .col("x", lance_datagen::array::step::()) + .into_reader_rows(RowCount::from(8), BatchCount::from(1)); + let dataset = Dataset::write(data, "memory://test_batch_readahead_concurrency", None) + .await + .unwrap(); + + // Default: threading mode falls back to get_num_compute_intensive_cpus(). + let plan = dataset.scan().create_plan().await.unwrap(); + let filtered = find_filtered_read(plan.as_ref()) + .expect("expected a FilteredReadExec in the scan plan"); + assert_eq!( + filtered.options().threading_mode, + FilteredReadThreadingMode::OnePartitionMultipleThreads(get_num_compute_intensive_cpus()), + ); + + // Explicit batch_readahead(N) bounds the decode fan-out to N. + let mut scanner = dataset.scan(); + scanner.batch_readahead(3); + let plan = scanner.create_plan().await.unwrap(); + let filtered = find_filtered_read(plan.as_ref()) + .expect("expected a FilteredReadExec in the scan plan"); + assert_eq!( + filtered.options().threading_mode, + FilteredReadThreadingMode::OnePartitionMultipleThreads(3), + ); + } + // The env var key scopes serial_test's lock so this test only blocks others // that touch LANCE_DEFAULT_IO_BUFFER_SIZE — unrelated tests still run in // parallel. diff --git a/rust/lance/src/io/exec/filtered_read.rs b/rust/lance/src/io/exec/filtered_read.rs index 059fa4e2b36..6c71a457b2b 100644 --- a/rust/lance/src/io/exec/filtered_read.rs +++ b/rust/lance/src/io/exec/filtered_read.rs @@ -360,11 +360,16 @@ impl FilteredReadStream { .clone() .unwrap_or_else(|| dataset.fragments().clone()); + let decode_parallelism = match threading_mode { + FilteredReadThreadingMode::OnePartitionMultipleThreads(n) + | FilteredReadThreadingMode::MultiplePartitions(n) => n, + }; log::debug!( - "Filtered read on {} fragments with frag_readahead={} and io_parallelism={}", + "Filtered read on {} fragments with frag_readahead={}, io_parallelism={} and decode_parallelism={}", fragments.len(), fragment_readahead, - io_parallelism + io_parallelism, + decode_parallelism ); // Ideally we don't need to collect here but if we don't we get "implementation of FnOnce is @@ -1482,6 +1487,16 @@ impl FilteredReadOptions { self.only_indexed_fragments = true; self } + + /// Specify the threading mode to use for the scan. + /// + /// This controls how decode work is parallelized. For the default single-partition + /// scan, the parameter of [`FilteredReadThreadingMode::OnePartitionMultipleThreads`] + /// bounds how many batch-decode tasks are buffered in flight (via `try_buffered`). + pub fn with_threading_mode(mut self, threading_mode: FilteredReadThreadingMode) -> Self { + self.threading_mode = threading_mode; + self + } } /// A plan node that reads a dataset, applying an optional filter and projection.