perf(remote): size blob file handles from one descriptor take (#4349)

python/python/tests/test_remote_db.py

@@ -2391,6 +2391,27 @@ def blob_query_response_table():
     )
 
 
+def blob_descriptor_take_table():
+    """What `fetch_blob_files` reads for rows 10, 20, 30, and 40."""
+    queried = blob_query_response_table()
+    image_field = queried.schema.field("image")
+    images = pa.StructArray.from_arrays(
+        [
+            pa.array([0, 0, 0, 0], type=pa.uint8()),
+            pa.array([0, 0, 0, 0], type=pa.uint64()),
+            pa.array([5, 0, 5, 0], type=pa.uint64()),
+            pa.array([0, 0, 0, 0], type=pa.uint32()),
+            pa.array(["", "", "", ""], type=pa.string()),
+        ],
+        fields=image_field.type,
+        mask=pa.array([False, True, False, False]),
+    )
+    return pa.Table.from_arrays(
+        [images, pa.array([10, 20, 30, 40], type=pa.uint64())],
+        schema=pa.schema([image_field, pa.field("_rowid", pa.uint64())]),
+    )
+
+
 @contextlib.contextmanager
 def blob_remote_table(*, server_version=Version("0.5.0")):
     def handler(request):
@@ -2429,9 +2450,14 @@ def blob_remote_table(*, server_version=Version("0.5.0")):
         elif request.path == "/v1/table/test/query/":
             content_len = int(request.headers.get("Content-Length", 0))
             body = json.loads(request.rfile.read(content_len))
-            assert body["columns"] == ["id", "image"]
             assert body["with_row_id"] is True
-            response_table = blob_query_response_table()
+            if body["columns"] == ["image"]:
+                # fetch_blob_files sizes its handles from a descriptor take.
+                assert body["filter"].startswith("_rowid IN")
+                response_table = blob_descriptor_take_table()
+            else:
+                assert body["columns"] == ["id", "image"]
+                response_table = blob_query_response_table()
             request.send_response(200)
             request.send_header("Content-Type", "application/vnd.apache.arrow.file")
             request.end_headers()

rust/lancedb/src/remote/table/blobs.rs

@@ -3,23 +3,28 @@
 
 //! Cloud blob column listing, whole-byte fetch, and seekable HTTP byte-range handles.
 
+use std::collections::HashMap;
 use std::ops::Range;
 use std::sync::Arc;
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::time::Duration;
 
-use arrow_array::{Array, LargeBinaryArray};
+use arrow_array::{Array, LargeBinaryArray, StructArray, UInt8Array, UInt64Array};
 use arrow_schema::DataType;
 use bytes::{Bytes, BytesMut};
+use datafusion_expr::{col, lit};
 use futures::{StreamExt, TryStreamExt};
+use lance::dataset::ROW_ID;
+use lance_core::datatypes::BlobKind;
 use reqwest::{Response, StatusCode, header};
 use tokio::sync::Mutex;
 
 use crate::Error;
 use crate::blob::BlobFile;
 use crate::error::Result;
+use crate::query::{QueryFilter, QueryRequest, Select};
 use crate::remote::client::{HttpSend, RequestResultExt, RestfulLanceDbClient};
-use crate::table::BaseTable;
+use crate::table::{AnyQuery, BaseTable};
 
 use super::{
     FreshnessHeaders, FreshnessState, ReadSnapshot, RemoteTable, VERSION_HEADER,
@@ -520,7 +525,11 @@ impl<S: HttpSend> RemoteTable<S> {
 
     /// Open seekable handles for `row_ids`.
     ///
-    /// Each non-null row is probed once to determine its size.
+    /// Descriptors for the requested rows come from one `_rowid` take query per
+    /// [`MAX_FETCH_BLOBS_ROW_IDS`] unique row ids, all read at one exact table
+    /// version. Each handle is sized from its descriptor and reads that
+    /// version. Only external blobs whose descriptor records no size are
+    /// probed, because their length is the whole object's.
     pub(super) async fn fetch_blob_files_impl(
         &self,
         column: &str,
@@ -530,46 +539,251 @@ impl<S: HttpSend> RemoteTable<S> {
         if row_ids.is_empty() {
             return Ok(Vec::new());
         }
-        if !self.server_version.support_blobs() {
-            return Err(Error::NotSupported {
-                message: "fetch_blob_files requires LanceDB Cloud server 0.5.0 or newer".into(),
-            });
-        }
+        self.ensure_blob_files_supported()?;
 
-        let read_snapshot = self.snapshot_read_state().await;
-        self.fetch_blob_files_with_snapshot(column, row_ids, read_snapshot)
-            .await
+        let mut read_snapshot = self.snapshot_read_state().await;
+        let sizes = self
+            .take_blob_descriptor_sizes(column, row_ids, &mut read_snapshot)
+            .await?;
+        let version = read_snapshot.version;
+
+        let mut files = Vec::with_capacity(row_ids.len());
+        let mut probe_indices = Vec::new();
+        let mut probe_requesters = Vec::new();
+        for (index, row_id) in row_ids.iter().enumerate() {
+            match sizes[row_id] {
+                DescriptorSize::Null => files.push(None),
+                DescriptorSize::Known(size) => {
+                    let requester = self.blob_range_requester(column, *row_id, &read_snapshot);
+                    files.push(Some(RemoteBlobFile::new(requester, size, version).into()));
+                }
+                DescriptorSize::Unresolved => {
+                    files.push(None);
+                    probe_indices.push(index);
+                    probe_requesters.push(self.blob_range_requester(
+                        column,
+                        *row_id,
+                        &read_snapshot,
+                    ));
+                }
+            }
+        }
+        let probed = probe_blob_files(probe_requesters).await?;
+        for (index, file) in probe_indices.into_iter().zip(probed) {
+            files[index] = file;
+        }
+        Ok(files)
     }
 
+    /// Probe each row once for its size.
+    ///
+    /// Serves callers that must not issue a query, such as a single blob that
+    /// the whole-byte route could not return.
     async fn fetch_blob_files_with_snapshot(
         &self,
         column: &str,
         row_ids: &[u64],
         read_snapshot: ReadSnapshot,
     ) -> Result<Vec<Option<BlobFile>>> {
-        let encoded_column = urlencoding::encode(column);
         let requesters = row_ids
             .iter()
-            .map(|row_id| {
-                let path = format!(
-                    "/v1/table/{}/blob/{encoded_column}/{row_id}/bytes",
-                    self.identifier
-                );
-                let requester: Arc<dyn BlobRangeRequester> = Arc::new(TableBlobRangeRequester {
-                    client: self.client.clone(),
-                    path,
-                    version: read_snapshot.version,
-                    branch: self.branch.clone(),
-                    freshness: Arc::new(std::sync::Mutex::new(read_snapshot.freshness_state)),
-                    parent_freshness: self.freshness.clone(),
-                    parent_freshness_request: read_snapshot.freshness,
-                    read_consistency_interval: self.client.read_consistency_interval,
-                });
-                requester
-            })
+            .map(|row_id| self.blob_range_requester(column, *row_id, &read_snapshot))
             .collect();
         probe_blob_files(requesters).await
     }
+
+    /// Read the blob descriptor of every unique row in `row_ids`.
+    ///
+    /// Pins `read_snapshot` to the version the first take read, so later chunks
+    /// and the returned handles all see one table version.
+    async fn take_blob_descriptor_sizes(
+        &self,
+        column: &str,
+        row_ids: &[u64],
+        read_snapshot: &mut ReadSnapshot,
+    ) -> Result<HashMap<u64, DescriptorSize>> {
+        let mut unique_row_ids = row_ids.to_vec();
+        unique_row_ids.sort_unstable();
+        unique_row_ids.dedup();
+
+        let mut sizes = HashMap::with_capacity(unique_row_ids.len());
+        for chunk in unique_row_ids.chunks(MAX_FETCH_BLOBS_ROW_IDS) {
+            let (mut chunk_sizes, response_version) = self
+                .take_blob_descriptor_chunk(column, chunk, *read_snapshot)
+                .await?;
+            if read_snapshot.version.is_none() {
+                match response_version {
+                    Some(version) => read_snapshot.version = Some(version),
+                    None => {
+                        // A server that does not report its read version may have read
+                        // an older version than a describe now returns, so the take is
+                        // repeated at the described version.
+                        read_snapshot.version =
+                            Some(self.describe_read_snapshot(*read_snapshot).await?.version);
+                        (chunk_sizes, _) = self
+                            .take_blob_descriptor_chunk(column, chunk, *read_snapshot)
+                            .await?;
+                    }
+                }
+            }
+            sizes.extend(chunk_sizes);
+        }
+        let found = unique_row_ids
+            .iter()
+            .filter(|row_id| sizes.contains_key(row_id))
+            .count();
+        if found < unique_row_ids.len() {
+            return Err(Error::InvalidInput {
+                message: format!(
+                    "blob read for column '{column}' requested {} row ids but only {found} exist \
+                     in the table; pass row ids collected from this table",
+                    unique_row_ids.len()
+                ),
+            });
+        }
+        Ok(sizes)
+    }
+
+    /// Take one chunk of descriptors, also returning the version the server
+    /// reported reading.
+    async fn take_blob_descriptor_chunk(
+        &self,
+        column: &str,
+        row_ids: &[u64],
+        read_snapshot: ReadSnapshot,
+    ) -> Result<(HashMap<u64, DescriptorSize>, Option<u64>)> {
+        let query = AnyQuery::Query(QueryRequest {
+            filter: Some(QueryFilter::Datafusion(
+                col(ROW_ID).in_list(row_ids.iter().map(|id| lit(*id)).collect(), false),
+            )),
+            select: Select::columns(&[column]),
+            with_row_id: true,
+            // Row ids name base-table rows, and the MemWAL scanner has no stable `_rowid`.
+            use_lsm: Some(false),
+            ..Default::default()
+        });
+        let body = self
+            .prepare_query_bodies(&query, read_snapshot.version)?
+            .pop()
+            .expect("a plain query has one body");
+        let request = self
+            .client
+            .post(&format!("/v1/table/{}/query/", self.identifier))
+            .json(&body);
+        let (request_id, response) = self
+            .send_with_freshness(request, true, read_snapshot.freshness)
+            .await?;
+        let response_version = response
+            .headers()
+            .get(&VERSION_HEADER)
+            .and_then(|value| value.to_str().ok())
+            .and_then(|value| value.parse().ok());
+        let mut stream = self.read_arrow_response(&request_id, response).await?;
+
+        let mut sizes = HashMap::with_capacity(row_ids.len());
+        let invalid_response = |message: String| Error::Http {
+            source: message.into(),
+            request_id: request_id.clone(),
+            status_code: None,
+        };
+        while let Some(batch) = stream.try_next().await? {
+            let taken_row_ids = batch
+                .column_by_name(ROW_ID)
+                .and_then(|array| array.as_any().downcast_ref::<UInt64Array>())
+                .ok_or_else(|| {
+                    invalid_response(format!("blob descriptor take is missing UInt64 '{ROW_ID}'"))
+                })?;
+            let values = batch.column_by_name(column).ok_or_else(|| {
+                invalid_response(format!("blob descriptor take is missing '{column}'"))
+            })?;
+            // Only blob v2 columns read back as descriptors.
+            let descriptors = values
+                .as_any()
+                .downcast_ref::<StructArray>()
+                .and_then(BlobDescriptorSizes::try_new)
+                .ok_or_else(|| Error::InvalidInput {
+                    message: format!("column '{column}' is not a blob column"),
+                })?;
+            for (index, row_id) in taken_row_ids.values().iter().enumerate() {
+                sizes.insert(*row_id, descriptors.size(index)?);
+            }
+        }
+        Ok((sizes, response_version))
+    }
+
+    fn ensure_blob_files_supported(&self) -> Result<()> {
+        if self.server_version.support_blobs() {
+            Ok(())
+        } else {
+            Err(Error::NotSupported {
+                message: "fetch_blob_files requires LanceDB Cloud server 0.5.0 or newer".into(),
+            })
+        }
+    }
+
+    fn blob_range_requester(
+        &self,
+        column: &str,
+        row_id: u64,
+        read_snapshot: &ReadSnapshot,
+    ) -> Arc<dyn BlobRangeRequester> {
+        let path = format!(
+            "/v1/table/{}/blob/{}/{row_id}/bytes",
+            self.identifier,
+            urlencoding::encode(column)
+        );
+        Arc::new(TableBlobRangeRequester {
+            client: self.client.clone(),
+            path,
+            version: read_snapshot.version,
+            branch: self.branch.clone(),
+            freshness: Arc::new(std::sync::Mutex::new(read_snapshot.freshness_state)),
+            parent_freshness: self.freshness.clone(),
+            parent_freshness_request: read_snapshot.freshness,
+            read_consistency_interval: self.client.read_consistency_interval,
+        })
+    }
+}
+
+/// What a blob v2 descriptor says about a row's size.
+#[derive(Debug, Clone, Copy)]
+enum DescriptorSize {
+    Null,
+    Known(u64),
+    /// An external blob that recorded no size. Its length is the whole
+    /// object's, which only the server can resolve.
+    Unresolved,
+}
+
+/// The `kind` and `size` children of a blob v2 descriptor array.
+struct BlobDescriptorSizes<'a> {
+    descriptors: &'a StructArray,
+    kinds: &'a UInt8Array,
+    sizes: &'a UInt64Array,
+}
+
+impl<'a> BlobDescriptorSizes<'a> {
+    fn try_new(descriptors: &'a StructArray) -> Option<Self> {
+        let child = |name: &str| descriptors.column_by_name(name);
+        Some(Self {
+            descriptors,
+            kinds: child("kind")?.as_any().downcast_ref::<UInt8Array>()?,
+            sizes: child("size")?.as_any().downcast_ref::<UInt64Array>()?,
+        })
+    }
+
+    fn size(&self, index: usize) -> Result<DescriptorSize> {
+        if self.descriptors.is_null(index) || self.kinds.is_null(index) {
+            return Ok(DescriptorSize::Null);
+        }
+        let kind = BlobKind::try_from(self.kinds.value(index))?;
+        let size = self.sizes.value(index);
+        Ok(if kind == BlobKind::External && size == 0 {
+            DescriptorSize::Unresolved
+        } else {
+            DescriptorSize::Known(size)
+        })
+    }
 }
 
 /// Probe blob sizes while preserving row order.
@@ -689,6 +903,8 @@ mod tests {
     use reqwest::Request;
     use semver::Version;
 
+    use arrow_array::RecordBatch;
+
     use super::*;
 
     const PAYLOAD: &[u8] = b"0123456789abcdefghijklmnopqrstuvwxyz";
@@ -772,10 +988,7 @@ mod tests {
         let requests = Arc::new(StdMutex::new(Vec::new()));
         let table = mock_remote_blob_table(requests.clone());
 
-        let mut files = table
-            .fetch_blob_files_impl("image", &[10, 20])
-            .await
-            .unwrap();
+        let mut files = table.probe_files("image", &[10, 20]).await.unwrap();
 
         assert_eq!(files.len(), 2);
         assert!(files[1].is_none());
@@ -787,12 +1000,368 @@ mod tests {
         );
     }
 
+    /// Descriptors as a remote query returns them, one `(kind, size)` per row.
+    fn descriptors(rows: &[Option<(BlobKind, u64)>]) -> StructArray {
+        let kinds = rows
+            .iter()
+            .map(|row| row.map_or(0, |(kind, _)| kind as u8))
+            .collect::<UInt8Array>();
+        let sizes = rows
+            .iter()
+            .map(|row| row.map_or(0, |(_, size)| size))
+            .collect::<UInt64Array>();
+        let len = rows.len();
+        StructArray::new(
+            lance_core::datatypes::BLOB_V2_DESC_FIELDS.clone(),
+            vec![
+                Arc::new(kinds),
+                Arc::new(UInt64Array::from(vec![0; len])),
+                Arc::new(sizes),
+                Arc::new(arrow_array::UInt32Array::from(vec![0; len])),
+                Arc::new(arrow_array::StringArray::from(vec![""; len])),
+            ],
+            Some(rows.iter().map(Option::is_some).collect()),
+        )
+    }
+
+    impl<S: HttpSend> RemoteTable<S> {
+        /// Open handles through the per-row size probe alone.
+        async fn probe_files(
+            &self,
+            column: &str,
+            row_ids: &[u64],
+        ) -> Result<Vec<Option<BlobFile>>> {
+            let read_snapshot = self.snapshot_read_state().await;
+            self.fetch_blob_files_with_snapshot(column, row_ids, read_snapshot)
+                .await
+        }
+    }
+
+    /// A request the take mock received.
+    #[derive(Debug, Clone, PartialEq)]
+    enum TakeMockRequest {
+        Query {
+            row_ids: Vec<u64>,
+            version: Option<u64>,
+        },
+        Describe,
+        Range {
+            row_id: u64,
+            range: String,
+            version: Option<String>,
+        },
+    }
+
+    /// Row ids named by a `_rowid IN (...)` filter.
+    fn filter_row_ids(filter: &str) -> Vec<u64> {
+        let list = filter.split_once('(').unwrap().1.trim_end_matches(')');
+        list.split(',')
+            .map(|id| id.trim().parse().unwrap())
+            .collect()
+    }
+
+    /// A table whose query route answers `_rowid` takes from `rows` and whose
+    /// byte routes serve [`PAYLOAD`]. Takes return rows in reverse order, and
+    /// report `query_version` in `x-lancedb-version` when it is set.
+    fn mock_take_blob_table(
+        rows: Vec<(u64, Option<(BlobKind, u64)>)>,
+        query_version: Option<&'static str>,
+    ) -> (
+        RemoteTable<crate::remote::client::test_utils::MockSender>,
+        Arc<StdMutex<Vec<TakeMockRequest>>>,
+    ) {
+        let requests = Arc::new(StdMutex::new(Vec::new()));
+        let captured = requests.clone();
+        let table = RemoteTable::new_mock(
+            "my_table".to_string(),
+            move |request| {
+                let path = request.url().path().to_string();
+                if path == "/v1/table/my_table/query/" {
+                    let body: serde_json::Value =
+                        serde_json::from_slice(request.body().unwrap().as_bytes().unwrap())
+                            .unwrap();
+                    assert_eq!(body["columns"], serde_json::json!(["image"]));
+                    assert_eq!(body["with_row_id"], serde_json::json!(true));
+                    assert_eq!(body["use_lsm"], serde_json::json!(false));
+                    let requested = filter_row_ids(body["filter"].as_str().unwrap());
+                    captured.lock().unwrap().push(TakeMockRequest::Query {
+                        row_ids: requested.clone(),
+                        version: body["version"].as_u64(),
+                    });
+                    let taken = rows
+                        .iter()
+                        .rev()
+                        .filter(|(row_id, _)| requested.contains(row_id))
+                        .collect::<Vec<_>>();
+                    let batch = RecordBatch::try_from_iter([
+                        (
+                            "image",
+                            Arc::new(descriptors(
+                                &taken.iter().map(|(_, row)| *row).collect::<Vec<_>>(),
+                            )) as Arc<dyn Array>,
+                        ),
+                        (
+                            ROW_ID,
+                            Arc::new(UInt64Array::from_iter_values(
+                                taken.iter().map(|(row_id, _)| *row_id),
+                            )) as Arc<dyn Array>,
+                        ),
+                    ])
+                    .unwrap();
+                    let mut body = Vec::new();
+                    let mut writer =
+                        arrow_ipc::writer::FileWriter::try_new(&mut body, &batch.schema()).unwrap();
+                    writer.write(&batch).unwrap();
+                    writer.finish().unwrap();
+                    drop(writer);
+                    let mut response = http::Response::builder().status(200);
+                    if let Some(version) = query_version {
+                        response = response.header(VERSION_HEADER, version);
+                    }
+                    return response.body(body).unwrap();
+                }
+                if path == "/v1/table/my_table/describe/" {
+                    captured.lock().unwrap().push(TakeMockRequest::Describe);
+                    return http::Response::builder()
+                        .status(200)
+                        .body(r#"{"version":9,"schema":{"fields":[]}}"#.as_bytes().to_vec())
+                        .unwrap();
+                }
+                let row_id = path
+                    .strip_prefix("/v1/table/my_table/blob/image/")
+                    .and_then(|rest| rest.strip_suffix("/bytes"))
+                    .unwrap_or_else(|| panic!("unexpected path: {path}"))
+                    .parse()
+                    .unwrap();
+                captured.lock().unwrap().push(TakeMockRequest::Range {
+                    row_id,
+                    range: request.headers()[header::RANGE]
+                        .to_str()
+                        .unwrap()
+                        .to_string(),
+                    version: request
+                        .url()
+                        .query_pairs()
+                        .find(|(name, _)| name == "version")
+                        .map(|(_, value)| value.into_owned()),
+                });
+                range_response(&request, PAYLOAD)
+            },
+            Some(Version::new(0, 5, 0)),
+        );
+        (table, requests)
+    }
+
+    #[tokio::test]
+    async fn remote_blob_files_are_sized_from_one_descriptor_take() {
+        let payload_size = PAYLOAD.len() as u64;
+        let (table, requests) = mock_take_blob_table(
+            vec![
+                (10, Some((BlobKind::Inline, payload_size))),
+                (20, None),
+                (30, Some((BlobKind::Inline, 0))),
+                (40, Some((BlobKind::Packed, payload_size))),
+                (50, Some((BlobKind::Dedicated, payload_size))),
+                (60, Some((BlobKind::External, 0))),
+            ],
+            Some("7"),
+        );
+
+        // Reordered and duplicated row ids come back in request order.
+        let files = table
+            .fetch_blob_files_impl("image", &[60, 10, 20, 30, 40, 50, 10])
+            .await
+            .unwrap();
+
+        let sizes = files
+            .iter()
+            .map(|file| file.as_ref().map(BlobFile::size))
+            .collect::<Vec<_>>();
+        assert_eq!(
+            sizes,
+            [
+                Some(payload_size),
+                Some(payload_size),
+                None,
+                Some(0),
+                Some(payload_size),
+                Some(payload_size),
+                Some(payload_size),
+            ]
+        );
+        // One take, then a probe for the external blob that recorded no size.
+        assert_eq!(
+            requests.lock().unwrap().as_slice(),
+            [
+                TakeMockRequest::Query {
+                    row_ids: vec![10, 20, 30, 40, 50, 60],
+                    version: None,
+                },
+                TakeMockRequest::Range {
+                    row_id: 60,
+                    range: "bytes=0-0".into(),
+                    version: Some("7".into()),
+                },
+            ]
+        );
+
+        // Sized handles read the version the take read.
+        assert_eq!(
+            files[1].as_ref().unwrap().read_range(5..12).await.unwrap(),
+            &PAYLOAD[5..12]
+        );
+        assert!(files[3].as_ref().unwrap().read().await.unwrap().is_empty());
+        assert_eq!(
+            requests.lock().unwrap().last().unwrap(),
+            &TakeMockRequest::Range {
+                row_id: 10,
+                range: "bytes=5-11".into(),
+                version: Some("7".into()),
+            }
+        );
+    }
+
+    #[tokio::test]
+    async fn remote_blob_files_take_every_chunk_at_the_first_version() {
+        let row_ids = (0..MAX_FETCH_BLOBS_ROW_IDS as u64 + 1).collect::<Vec<_>>();
+        let (table, requests) = mock_take_blob_table(
+            row_ids.iter().map(|row_id| (*row_id, None)).collect(),
+            Some("7"),
+        );
+
+        let files = table
+            .fetch_blob_files_impl("image", &row_ids)
+            .await
+            .unwrap();
+
+        assert_eq!(files.len(), row_ids.len());
+        let versions = requests
+            .lock()
+            .unwrap()
+            .iter()
+            .map(|request| match request {
+                TakeMockRequest::Query { row_ids, version } => (row_ids.len(), *version),
+                other => panic!("unexpected request: {other:?}"),
+            })
+            .collect::<Vec<_>>();
+        assert_eq!(versions, [(MAX_FETCH_BLOBS_ROW_IDS, None), (1, Some(7))]);
+    }
+
+    #[tokio::test]
+    async fn remote_blob_files_retake_at_the_described_version_when_a_take_does_not_report_it() {
+        let payload_size = PAYLOAD.len() as u64;
+        let (table, requests) =
+            mock_take_blob_table(vec![(10, Some((BlobKind::Inline, payload_size)))], None);
+
+        let file = table
+            .fetch_blob_files_impl("image", &[10])
+            .await
+            .unwrap()
+            .pop()
+            .flatten()
+            .unwrap();
+        file.read_range(0..4).await.unwrap();
+
+        // The unversioned take is repeated at the described version.
+        assert_eq!(
+            requests.lock().unwrap().as_slice(),
+            [
+                TakeMockRequest::Query {
+                    row_ids: vec![10],
+                    version: None,
+                },
+                TakeMockRequest::Describe,
+                TakeMockRequest::Query {
+                    row_ids: vec![10],
+                    version: Some(9),
+                },
+                TakeMockRequest::Range {
+                    row_id: 10,
+                    range: "bytes=0-3".into(),
+                    version: Some("9".into()),
+                },
+            ]
+        );
+    }
+
+    #[tokio::test]
+    async fn remote_blob_files_take_at_the_checked_out_version() {
+        let payload_size = PAYLOAD.len() as u64;
+        let (table, requests) = mock_take_blob_table(
+            vec![(10, Some((BlobKind::Inline, payload_size)))],
+            Some("7"),
+        );
+        table.checkout(9).await.unwrap();
+        requests.lock().unwrap().clear();
+
+        table.fetch_blob_files_impl("image", &[10]).await.unwrap();
+
+        assert_eq!(
+            requests.lock().unwrap().as_slice(),
+            [TakeMockRequest::Query {
+                row_ids: vec![10],
+                version: Some(9),
+            }]
+        );
+    }
+
+    #[tokio::test]
+    async fn remote_blob_files_reject_row_ids_missing_from_the_take() {
+        let (table, requests) = mock_take_blob_table(vec![(10, None)], Some("7"));
+
+        let error = table
+            .fetch_blob_files_impl("image", &[10, 99])
+            .await
+            .unwrap_err();
+
+        assert!(matches!(error, Error::InvalidInput { .. }), "{error}");
+        assert!(error.to_string().contains("only 1 exist"), "{error}");
+        assert_eq!(requests.lock().unwrap().len(), 1);
+    }
+
+    #[tokio::test]
+    async fn remote_blob_files_reject_a_non_blob_column() {
+        let table = RemoteTable::new_mock(
+            "my_table".to_string(),
+            |request| {
+                assert_eq!(request.url().path(), "/v1/table/my_table/query/");
+                let batch = RecordBatch::try_from_iter([
+                    (
+                        "id",
+                        Arc::new(arrow_array::Int64Array::from(vec![1])) as Arc<dyn Array>,
+                    ),
+                    (
+                        ROW_ID,
+                        Arc::new(UInt64Array::from(vec![10])) as Arc<dyn Array>,
+                    ),
+                ])
+                .unwrap();
+                let mut body = Vec::new();
+                let mut writer =
+                    arrow_ipc::writer::FileWriter::try_new(&mut body, &batch.schema()).unwrap();
+                writer.write(&batch).unwrap();
+                writer.finish().unwrap();
+                drop(writer);
+                http::Response::builder().status(200).body(body).unwrap()
+            },
+            Some(Version::new(0, 5, 0)),
+        );
+
+        let error = table.fetch_blob_files_impl("id", &[10]).await.unwrap_err();
+
+        assert!(matches!(error, Error::InvalidInput { .. }), "{error}");
+        assert!(
+            error.to_string().contains("'id' is not a blob column"),
+            "{error}"
+        );
+    }
+
     #[tokio::test]
     async fn remote_blob_file_reads_the_requested_range() {
         let requests = Arc::new(StdMutex::new(Vec::new()));
         let table = mock_remote_blob_table(requests.clone());
         let file = table
-            .fetch_blob_files_impl("image", &[10])
+            .probe_files("image", &[10])
             .await
             .unwrap()
             .pop()
@@ -841,7 +1410,7 @@ mod tests {
         );
 
         let file = table
-            .fetch_blob_files_impl("image", &[10])
+            .probe_files("image", &[10])
             .await
             .unwrap()
             .pop()
@@ -890,7 +1459,7 @@ mod tests {
 
         table.checkout(5).await.unwrap();
         let file = table
-            .fetch_blob_files_impl("image", &[10])
+            .probe_files("image", &[10])
             .await
             .unwrap()
             .pop()
@@ -910,7 +1479,7 @@ mod tests {
         let requests = Arc::new(StdMutex::new(Vec::new()));
         let table = mock_remote_blob_table(requests.clone());
         let file = table
-            .fetch_blob_files_impl("image", &[10])
+            .probe_files("image", &[10])
             .await
             .unwrap()
             .pop()
@@ -933,10 +1502,7 @@ mod tests {
         let requests = Arc::new(StdMutex::new(Vec::new()));
         let table = mock_remote_blob_table(requests.clone());
 
-        let mut files = table
-            .fetch_blob_files_impl("image", &[10, 20, 30])
-            .await
-            .unwrap();
+        let mut files = table.probe_files("image", &[10, 20, 30]).await.unwrap();
 
         assert_eq!(files.len(), 3);
         assert!(files[1].is_none());
@@ -964,10 +1530,7 @@ mod tests {
             Some(Version::new(0, 5, 0)),
         );
 
-        let error = table
-            .fetch_blob_files_impl("image", &[10])
-            .await
-            .unwrap_err();
+        let error = table.probe_files("image", &[10]).await.unwrap_err();
         assert!(
             error
                 .to_string()
@@ -991,10 +1554,7 @@ mod tests {
             Some(Version::new(0, 5, 0)),
         );
 
-        let error = table
-            .fetch_blob_files_impl("image", &[10])
-            .await
-            .unwrap_err();
+        let error = table.probe_files("image", &[10]).await.unwrap_err();
         assert!(
             error
                 .to_string()
@@ -1015,10 +1575,7 @@ mod tests {
             Some(Version::new(0, 5, 0)),
         );
 
-        let error = table
-            .fetch_blob_files_impl("image", &[10])
-            .await
-            .unwrap_err();
+        let error = table.probe_files("image", &[10]).await.unwrap_err();
         assert!(
             error
                 .to_string()
@@ -1076,7 +1633,7 @@ mod tests {
             Some(Version::new(0, 5, 0)),
         );
         let file = table
-            .fetch_blob_files_impl("image", &[10])
+            .probe_files("image", &[10])
             .await
             .unwrap()
             .pop()
@@ -1116,7 +1673,7 @@ mod tests {
             Some(Version::new(0, 5, 0)),
         );
         let file = table
-            .fetch_blob_files_impl("image", &[10])
+            .probe_files("image", &[10])
             .await
             .unwrap()
             .pop()
@@ -1167,7 +1724,7 @@ mod tests {
             Some(Version::new(0, 5, 0)),
         );
         let file = table
-            .fetch_blob_files_impl("image", &[10])
+            .probe_files("image", &[10])
             .await
             .unwrap()
             .pop()
@@ -1194,7 +1751,7 @@ mod tests {
         let requests = Arc::new(StdMutex::new(Vec::new()));
         let table = mock_remote_blob_table(requests.clone());
         let file = table
-            .fetch_blob_files_impl("image", &[10])
+            .probe_files("image", &[10])
             .await
             .unwrap()
             .pop()
@@ -1273,7 +1830,7 @@ mod tests {
         let requests = Arc::new(StdMutex::new(Vec::new()));
         let table = mock_remote_blob_table(requests.clone());
 
-        let mut files = table.fetch_blob_files_impl("image", &[10]).await.unwrap();
+        let mut files = table.probe_files("image", &[10]).await.unwrap();
         let file = files.remove(0).unwrap();
 
         assert_eq!(file.size(), PAYLOAD.len() as u64);
@@ -1288,7 +1845,7 @@ mod tests {
         let requests = Arc::new(StdMutex::new(Vec::new()));
         let table = mock_remote_blob_table(requests.clone());
 
-        let mut files = table.fetch_blob_files_impl("image", &[10]).await.unwrap();
+        let mut files = table.probe_files("image", &[10]).await.unwrap();
         let file = files.remove(0).unwrap();
         let probe_requests = requests.lock().unwrap().len();
 
@@ -1313,7 +1870,7 @@ mod tests {
         let requests = Arc::new(StdMutex::new(Vec::new()));
         let table = mock_remote_blob_table(requests.clone());
 
-        let mut files = table.fetch_blob_files_impl("image", &[10]).await.unwrap();
+        let mut files = table.probe_files("image", &[10]).await.unwrap();
         let file = files.remove(0).unwrap();
         let probe_requests = requests.lock().unwrap().len();
 
@@ -1365,7 +1922,7 @@ mod tests {
             Some(Version::new(0, 5, 0)),
         );
 
-        let mut files = table.fetch_blob_files_impl("image", &[10]).await.unwrap();
+        let mut files = table.probe_files("image", &[10]).await.unwrap();
         let file = files.remove(0).unwrap();
 
         let error = file.read_range(1..3).await.unwrap_err();
@@ -1456,7 +2013,7 @@ mod tests {
         let requests = Arc::new(StdMutex::new(Vec::new()));
         let table = mock_remote_blob_table(requests.clone());
         let file = table
-            .fetch_blob_files_impl("image", &[10])
+            .probe_files("image", &[10])
             .await
             .unwrap()
             .pop()
@@ -1502,7 +2059,7 @@ mod tests {
             Some(Version::new(0, 5, 0)),
         );
 
-        let mut files = table.fetch_blob_files_impl("image", &[10]).await.unwrap();
+        let mut files = table.probe_files("image", &[10]).await.unwrap();
         let file = files.remove(0).unwrap();
         let error = file.read_range(1..3).await.unwrap_err();
         assert!(error.to_string().contains("206"), "got: {error}");