feat(python): preserve re-scannable merge sources

python/python/lance/dataset.py

@@ -78,7 +78,7 @@
 from .lance import __version__ as __version__
 from .lance import _Session as Session
 from .query import DocumentGranularity, FullTextQuery, MinHashQuery
-from .types import _coerce_reader, _is_materialized
+from .types import WriteSource, _coerce_reader, coerce_source
 from .udf import BatchUDF, normalize_transform
 from .udf import BatchUDFCheckpoint as BatchUDFCheckpoint
 from .udf import batch_udf as batch_udf
@@ -393,25 +393,32 @@ def execute(self, data_obj: ReaderLike, *, schema: Optional[pa.Schema] = None):
         schema: Optional[pa.Schema]
             The schema of the data.  This only needs to be supplied whenever the data
             source is some kind of generator.
-        """
-        reader = _coerce_reader(data_obj, schema)
 
-        # Materialized sources are wrapped in an in-memory table so retries never
-        # spill and the source's statistics can drive the join; everything else is
-        # treated as a one-shot stream.
-        if _is_materialized(data_obj):
+        Notes
+        -----
+        :func:`lance.types.coerce_source` adapts built-in and registered source
+        types. Repeatable sources provide fresh readers for conflict retries.
+        Keep external source files unchanged until the operation finishes.
+        Sources without a repeatability contract retain the replay buffer.
+        """
+        source = coerce_source(data_obj, schema)
+        if source.mode == "replayable":
+            return super(MergeInsertBuilder, self).execute_replayable(
+                source.schema, source.open_reader
+            )
+        reader = source.open_reader()
+        if source.mode == "materialized":
             return super(MergeInsertBuilder, self).execute_batches(reader)
-
         return super(MergeInsertBuilder, self).execute(reader)
 
     def execute_uncommitted(
         self, data_obj: ReaderLike, *, schema: Optional[pa.Schema] = None
     ) -> Tuple[Transaction, Dict[str, Any]]:
         """Executes the merge insert operation without committing
 
-        This function updates the original dataset and returns a dictionary with
-        information about merge statistics - i.e. the number of inserted, updated,
-        and deleted rows.
+        This function returns a transaction and merge statistics. It does not
+        commit the transaction or update the original dataset. The source is
+        consumed once, without a buffer for commit retries.
 
         Parameters
         ----------
@@ -424,9 +431,9 @@ def execute_uncommitted(
             The schema of the data.  This only needs to be supplied whenever the data
             source is some kind of generator.
         """
-        reader = _coerce_reader(data_obj, schema)
-
-        if _is_materialized(data_obj):
+        source = coerce_source(data_obj, schema)
+        reader = source.open_reader()
+        if source.mode == "materialized":
             return super(MergeInsertBuilder, self).execute_uncommitted_batches(reader)
 
         return super(MergeInsertBuilder, self).execute_uncommitted(reader)
@@ -819,12 +826,13 @@ def analyze_plan(
         - iops: number of I/O operations performed
         - requests: number of storage requests made
         """  # noqa: E501
-        reader = _coerce_reader(data_obj, schema)
-
-        # Route exactly as execute() does, so the reported plan is the one that
-        # would run. A materialized source reports exact statistics where a stream
-        # reports none, which can change which side of the join is collected.
-        if _is_materialized(data_obj):
+        source = coerce_source(data_obj, schema)
+        if source.mode == "replayable":
+            return super(MergeInsertBuilder, self).analyze_plan_replayable(
+                source.schema, source.open_reader
+            )
+        reader = source.open_reader()
+        if source.mode == "materialized":
             return super(MergeInsertBuilder, self).analyze_plan_batches(reader)
 
         return super(MergeInsertBuilder, self).analyze_plan(reader)
@@ -8831,3 +8839,50 @@ def __init__(
 
     def inner(self):
         return self._inner
+
+
+@coerce_source.register(LanceDataset)
+def _lance_dataset_source(data_obj: LanceDataset, schema=None) -> WriteSource:
+    # Capture the native dataset version and all default scan options once.
+    return coerce_source(data_obj.scanner(), schema)
+
+
+@coerce_source.register(LanceScanner)
+def _lance_scanner_source(data_obj: LanceScanner, schema=None) -> WriteSource:
+    if type(data_obj) is not LanceScanner:
+        return WriteSource.from_reader(data_obj.to_reader())
+    snapshot = data_obj._snapshot
+    replayable = False
+    if snapshot is not None:
+        defaults = vars(ScannerBuilder(data_obj._ds)).copy()
+        defaults.pop("ds")
+        # These settings preserve source values and order. Expressions, searches,
+        # callbacks, and sorting can change results across otherwise valid scans.
+        stable_options = {
+            "_columns",
+            "_limit",
+            "_offset",
+            "_batch_size",
+            "_batch_size_bytes",
+            "_io_buffer_size",
+            "_batch_readahead",
+            "_fragment_readahead",
+            "_fragments",
+            "_with_row_id",
+            "_with_row_address",
+            "_use_stats",
+            "_use_scalar_index",
+            "_include_deleted_rows",
+            "_late_materialization",
+            "_blob_handling",
+            "_strict_batch_size",
+            "_row_addr_allowlist",
+            "_row_addr_blocklist",
+        }
+        replayable = snapshot.keys() == defaults.keys() and all(
+            key in stable_options or value == defaults[key]
+            for key, value in snapshot.items()
+        )
+    if replayable:
+        return WriteSource(data_obj.projected_schema, data_obj.to_reader)
+    return WriteSource.from_reader(data_obj.to_reader())

python/python/lance/lance/__init__.pyi

@@ -719,6 +719,12 @@ class _MergeInsertBuilder:
     def target_all_bases(self, include_primary: bool = True) -> Self: ...
     def execute(self, new_data: pa.RecordBatchReader) -> ExecuteResult: ...
     def execute_batches(self, new_data: pa.RecordBatchReader) -> ExecuteResult: ...
+    def execute_replayable(
+        self, schema: pa.Schema, factory: Callable[[], pa.RecordBatchReader]
+    ) -> ExecuteResult: ...
+    def analyze_plan_replayable(
+        self, schema: pa.Schema, factory: Callable[[], pa.RecordBatchReader]
+    ) -> str: ...
     def execute_uncommitted(
         self, new_data: pa.RecordBatchReader
     ) -> tuple[Transaction, ExecuteResult]: ...