Support partial blob reads

adapters/python/coflux/context.py

@@ -338,12 +338,21 @@ def get_asset_entries(self, asset_id: str) -> list[AssetEntry]:
             entries.append(AssetEntry(path, blob_key, size, metadata or {}))
         return entries
 
-    def download_blob(self, blob_key: str, target_path: Path) -> None:
-        """Download a blob to a local file."""
+    def download_blob(
+        self,
+        blob_key: str,
+        target_path: Path,
+        *,
+        offset: int | None = None,
+        length: int | None = None,
+    ) -> None:
+        """Download a blob, or a byte range of one, to a local file."""
         request_id = protocol.request_download_blob(
             self.execution_id,
             blob_key,
             str(target_path),
+            offset,
+            length,
         )
         self._wait_response(request_id)
 
@@ -450,10 +459,10 @@ def create_asset(
         if not paths_to_upload and not resolved_entries:
             raise ValueError("No files found to create asset")
 
-        abs_paths = [str(p) for _, p in paths_to_upload] if paths_to_upload else None
+        upload_paths = {rel: str(p) for rel, p in paths_to_upload} or None
         request_id = protocol.request_persist_asset(
             self.execution_id,
-            abs_paths,
+            upload_paths,
             {"name": name} if name else None,
             resolved_entries if resolved_entries else None,
         )

adapters/python/coflux/models.py

@@ -4,6 +4,7 @@
 
 import fnmatch
 import functools
+import tempfile
 import typing as t
 from pathlib import Path
 
@@ -42,6 +43,21 @@ def restore(self, *, at: Path | str | None = None) -> Path:
         ctx.download_blob(self.blob_key, target)
         return target
 
+    def read(self, offset: int = 0, length: int | None = None) -> bytes:
+        """Read bytes from this entry without restoring the whole file.
+
+        Reads to the end of the file when ``length`` is omitted.
+
+        Useful for formats that seek rather than read straight through — a
+        Parquet footer, say — where restoring a large file to read a few
+        kilobytes of it would be wasteful.
+        """
+        ctx = get_context()
+        with tempfile.TemporaryDirectory() as directory:
+            target = Path(directory) / "range"
+            ctx.download_blob(self.blob_key, target, offset=offset, length=length)
+            return target.read_bytes()
+
 
 class AssetMetadata(t.NamedTuple):
     """Metadata for an asset reference."""

adapters/python/coflux/protocol.py

@@ -265,15 +265,15 @@ def request_select(
 
 def request_persist_asset(
     execution_id: str,
-    paths: list[str] | None = None,
+    paths: dict[str, str] | None = None,
     metadata: dict[str, Any] | None = None,
     entries: dict[str, tuple[str, int, dict[str, Any]]] | None = None,
 ) -> int:
     """Request to persist an asset.
 
     Args:
         execution_id: The execution this asset belongs to.
-        paths: Local file paths to upload and include.
+        paths: Local files to upload, as {path within the asset: local path}.
         metadata: Asset-level metadata (e.g. name).
         entries: Pre-resolved entries referencing existing blobs.
             Each value is (blob_key, size, entry_metadata).
@@ -311,16 +311,22 @@ def request_download_blob(
     execution_id: str,
     blob_key: str,
     target_path: str,
+    offset: int | None = None,
+    length: int | None = None,
 ) -> int:
-    """Request to download a blob to a local file."""
-    return get_protocol().send_request(
-        "download_blob",
-        {
-            "execution_id": execution_id,
-            "blob_key": blob_key,
-            "target_path": target_path,
-        },
-    )
+    """Request to download a blob, or a byte range of one, to a local file."""
+    params: dict[str, Any] = {
+        "execution_id": execution_id,
+        "blob_key": blob_key,
+        "target_path": target_path,
+    }
+    # Omitted rather than sent as null, so the message keeps the shape
+    # older CLIs expect.
+    if offset is not None:
+        params["offset"] = offset
+    if length is not None:
+        params["length"] = length
+    return get_protocol().send_request("download_blob", params)
 
 
 def request_upload_blob(

cli/internal/adapter/protocol.go

@@ -234,10 +234,10 @@ type SelectParams struct {
 
 // PersistAssetParams for persist_asset request
 type PersistAssetParams struct {
-	ExecutionID string           `json:"execution_id"`
-	Paths       []string         `json:"paths,omitempty"`
-	Metadata    map[string]any   `json:"metadata,omitempty"`
-	Entries     map[string][]any `json:"entries,omitempty"` // Pre-resolved entries: {path: [blob_key, size, metadata]}
+	ExecutionID string            `json:"execution_id"`
+	Paths       map[string]string `json:"paths,omitempty"` // Files to upload: {path within the asset: local path}
+	Metadata    map[string]any    `json:"metadata,omitempty"`
+	Entries     map[string][]any  `json:"entries,omitempty"` // Pre-resolved entries: {path: [blob_key, size, metadata]}
 }
 
 // PersistAssetResult is the response to persist_asset
@@ -434,11 +434,14 @@ type StreamClosedParams struct {
 	Error          map[string]any `json:"error,omitempty"`
 }
 
-// DownloadBlobParams for download_blob request
+// DownloadBlobParams for download_blob request. Offset and length are
+// optional: with neither, the whole blob is downloaded.
 type DownloadBlobParams struct {
 	ExecutionID string `json:"execution_id"`
 	BlobKey     string `json:"blob_key"`
 	TargetPath  string `json:"target_path"`
+	Offset      *int64 `json:"offset,omitempty"`
+	Length      *int64 `json:"length,omitempty"`
 }
 
 // UploadBlobParams for upload_blob request

cli/internal/blob/blob.go

@@ -14,6 +14,11 @@ import (
 type Store interface {
 	// Get retrieves a blob by key, returns nil if not found
 	Get(key string) (io.ReadCloser, error)
+	// GetRange retrieves a byte range of a blob, returns nil if not found.
+	// A negative length reads to the end of the blob.
+	GetRange(key string, offset, length int64) (io.ReadCloser, error)
+	// Exists reports whether a blob is already stored
+	Exists(key string) (bool, error)
 	// Put stores a blob and returns its key (content-addressed)
 	Put(reader io.Reader) (string, error)
 	// Upload uploads a file and returns its key
@@ -22,6 +27,34 @@ type Store interface {
 	Download(key, path string) (bool, error)
 }
 
+// Blobs at or above this size are checked for existence before being
+// uploaded.
+const existsCheckThreshold = 1 << 20 // 1 MiB
+
+// skipUpload reports whether content with this key is already stored, and
+// so needn't be sent again. Content addressing makes re-uploading it
+// redundant, but asking costs a round trip, so it's only worth it once the
+// content is big enough that re-sending would cost more than asking.
+//
+// A failed check reports "not stored": this is only an optimisation, and
+// uploading something that's already there is always safe.
+func skipUpload(store Store, key string, size int) bool {
+	if size < existsCheckThreshold {
+		return false
+	}
+	exists, err := store.Exists(key)
+	return err == nil && exists
+}
+
+// rangeHeader formats an HTTP byte range. A negative length means "to the
+// end of the blob".
+func rangeHeader(offset, length int64) string {
+	if length < 0 {
+		return fmt.Sprintf("bytes=%d-", offset)
+	}
+	return fmt.Sprintf("bytes=%d-%d", offset, offset+length-1)
+}
+
 // Manager manages multiple blob stores with fallback
 type Manager struct {
 	stores    []Store
@@ -52,6 +85,24 @@ func (m *Manager) Get(key string) (io.ReadCloser, error) {
 	return nil, fmt.Errorf("blob not found: %s", key)
 }
 
+// GetRange retrieves a byte range of a blob from any store. A negative
+// length reads to the end of the blob.
+func (m *Manager) GetRange(key string, offset, length int64) (io.ReadCloser, error) {
+	// Deliberately not served from the cache: it holds whole blobs, and
+	// Download treats the file merely existing as a complete one, so a
+	// range must never be written there.
+	for _, store := range m.stores {
+		reader, err := store.GetRange(key, offset, length)
+		if err != nil {
+			return nil, err
+		}
+		if reader != nil {
+			return reader, nil
+		}
+	}
+	return nil, fmt.Errorf("blob not found: %s", key)
+}
+
 // Put stores a blob in the first store
 func (m *Manager) Put(reader io.Reader) (string, error) {
 	if len(m.stores) == 0 {
@@ -119,6 +170,33 @@ func (m *Manager) DownloadTo(key, targetPath string) error {
 	return fmt.Errorf("blob not found: %s", key)
 }
 
+// DownloadRangeTo downloads a byte range of a blob to a specific path. A
+// negative length reads to the end of the blob.
+func (m *Manager) DownloadRangeTo(key, targetPath string, offset, length int64) error {
+	if err := os.MkdirAll(filepath.Dir(targetPath), 0755); err != nil {
+		return err
+	}
+
+	f, err := os.Create(targetPath)
+	if err != nil {
+		return err
+	}
+	defer func() { _ = f.Close() }()
+
+	if length == 0 {
+		return nil
+	}
+
+	reader, err := m.GetRange(key, offset, length)
+	if err != nil {
+		return err
+	}
+	defer func() { _ = reader.Close() }()
+
+	_, err = io.Copy(f, reader)
+	return err
+}
+
 // CachePath returns the cache path for a blob key
 func (m *Manager) CachePath(key string) string {
 	// Use first 4 chars for directory sharding

cli/internal/blob/blob_test.go

@@ -0,0 +1,143 @@
+package blob
+
+import (
+	"bytes"
+	"errors"
+	"fmt"
+	"io"
+	"testing"
+)
+
+// fakeStore records what it was asked, so a test can tell an upload that
+// happened from one that was skipped.
+type fakeStore struct {
+	content     map[string][]byte
+	existsErr   error
+	existsCalls int
+}
+
+func (f *fakeStore) Get(key string) (io.ReadCloser, error) {
+	content, ok := f.content[key]
+	if !ok {
+		return nil, nil
+	}
+	return io.NopCloser(bytes.NewReader(content)), nil
+}
+
+func (f *fakeStore) GetRange(key string, offset, length int64) (io.ReadCloser, error) {
+	content, ok := f.content[key]
+	if !ok {
+		return nil, nil
+	}
+	if offset > int64(len(content)) {
+		return nil, fmt.Errorf("range not satisfiable")
+	}
+	end := int64(len(content))
+	if length >= 0 && offset+length < end {
+		end = offset + length
+	}
+	return io.NopCloser(bytes.NewReader(content[offset:end])), nil
+}
+
+func (f *fakeStore) Exists(key string) (bool, error) {
+	f.existsCalls++
+	if f.existsErr != nil {
+		return false, f.existsErr
+	}
+	_, ok := f.content[key]
+	return ok, nil
+}
+
+func (f *fakeStore) Put(reader io.Reader) (string, error)    { return "", nil }
+func (f *fakeStore) Upload(path string) (string, error)      { return "", nil }
+func (f *fakeStore) Download(key, path string) (bool, error) { return false, nil }
+
+func stored(key string, content string) *fakeStore {
+	return &fakeStore{content: map[string][]byte{key: []byte(content)}}
+}
+
+func TestSkipUploadBelowThreshold(t *testing.T) {
+	store := stored("abc", "hello")
+	if skipUpload(store, "abc", existsCheckThreshold-1) {
+		t.Fatal("expected small content to be uploaded without checking")
+	}
+	if store.existsCalls != 0 {
+		t.Fatalf("expected no existence check, got %d", store.existsCalls)
+	}
+}
+
+func TestSkipUploadWhenPresent(t *testing.T) {
+	if !skipUpload(stored("abc", "hello"), "abc", existsCheckThreshold) {
+		t.Fatal("expected stored content to be skipped")
+	}
+}
+
+func TestSkipUploadWhenAbsent(t *testing.T) {
+	if skipUpload(stored("abc", "hello"), "other", existsCheckThreshold) {
+		t.Fatal("expected absent content to be uploaded")
+	}
+}
+
+func TestSkipUploadOnCheckFailure(t *testing.T) {
+	// The check is an optimisation, so a store that can't answer must not
+	// stop the upload.
+	store := stored("abc", "hello")
+	store.existsErr = errors.New("nope")
+	if skipUpload(store, "abc", existsCheckThreshold) {
+		t.Fatal("expected a failed check to fall back to uploading")
+	}
+}
+
+func TestRangeHeader(t *testing.T) {
+	for _, c := range []struct {
+		offset, length int64
+		want           string
+	}{
+		{0, 16, "bytes=0-15"},
+		{1000, 16, "bytes=1000-1015"},
+		{1000, -1, "bytes=1000-"},
+	} {
+		if got := rangeHeader(c.offset, c.length); got != c.want {
+			t.Errorf("rangeHeader(%d, %d) = %q, want %q", c.offset, c.length, got, c.want)
+		}
+	}
+}
+
+func TestManagerGetRangeFallsThroughStores(t *testing.T) {
+	// The blob is missing from the first store, so the second answers.
+	m := NewManager([]Store{stored("other", "xxx"), stored("abc", "0123456789")}, t.TempDir(), 200)
+
+	reader, err := m.GetRange("abc", 3, 4)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() { _ = reader.Close() }()
+	got, err := io.ReadAll(reader)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if string(got) != "3456" {
+		t.Fatalf("got %q, want %q", got, "3456")
+	}
+}
+
+func TestManagerGetRangeToEnd(t *testing.T) {
+	m := NewManager([]Store{stored("abc", "0123456789")}, t.TempDir(), 200)
+
+	reader, err := m.GetRange("abc", 6, -1)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() { _ = reader.Close() }()
+	got, _ := io.ReadAll(reader)
+	if string(got) != "6789" {
+		t.Fatalf("got %q, want %q", got, "6789")
+	}
+}
+
+func TestManagerGetRangeNotFound(t *testing.T) {
+	m := NewManager([]Store{stored("abc", "0123456789")}, t.TempDir(), 200)
+	if _, err := m.GetRange("missing", 0, 4); err == nil {
+		t.Fatal("expected an error for a blob no store holds")
+	}
+}