A Deep Deep Dive into df.write_lance()

[everettVT]

So as a part of another discussion on adding LanceTable support to the daft catalog, I began reviewing the daft.read_lance and df.write_lance methods and went down a huge rabbit hole which helped me understand both how daft performs distributed writes, and how lance fragements work.

I figured I would share my notes that I made a long the way as an exercise.

How Does Daft write Lance Datasets?

sequenceDiagram
    participant User
    participant DataFrame
    participant WriteSink as DataFrame.write_sink
    participant LanceDataSink
    participant LogicalPlanBuilder
    participant RustLogicalBuilder as _LogicalPlanBuilder (Rust)
    participant PhysicalPlan
    participant DataSink as DataSink.write()
    participant Lance
    
    User->>DataFrame: write_lance(uri, mode, io_config, **kwargs)
    DataFrame->>LanceDataSink: new LanceDataSink(uri, schema, mode, io_config, **kwargs)
    
    Note over LanceDataSink: Constructor:<br/>- Set _table_uri, _mode, _io_config, _kwargs<br/>- Convert io_config to storage_options<br/>- Convert schema to pyarrow_schema<br/>- Try lance.dataset() to read existing<br/>- Set _version (0 or latest_version)<br/>- Validate schema if table exists
    
    LanceDataSink-->>DataFrame: sink instance
    DataFrame->>WriteSink: write_sink(sink)
    
    WriteSink->>LanceDataSink: sink.start()
    Note over LanceDataSink: start() is pass-through (no-op)
    
    WriteSink->>LogicalPlanBuilder: _builder.write_datasink(sink.name(), sink)
    LogicalPlanBuilder->>RustLogicalBuilder: _builder.datasink_write(name, sink)
    RustLogicalBuilder-->>LogicalPlanBuilder: new _LogicalPlanBuilder
    LogicalPlanBuilder-->>WriteSink: new LogicalPlanBuilder
    
    WriteSink->>WriteSink: write_df = DataFrame(builder)
    WriteSink->>WriteSink: write_df.collect()
    
    Note over WriteSink,PhysicalPlan: Physical execution begins during collect()
    WriteSink->>PhysicalPlan: Execute physical plan
    
    loop For each partition in the DataFrame
        PhysicalPlan->>DataSink: sink.write(micropartitions_iterator)
        
        loop For each micropartition in partition
            DataSink->>DataSink: micropartition.to_arrow()
            DataSink->>Lance: lance.fragment.write_fragments(arrow_table, dataset_uri, mode, storage_options, **kwargs)
            Lance-->>DataSink: list[FragmentMetadata]
            DataSink->>DataSink: yield WriteResult(fragments, bytes_written, rows_written)
        end
        
        DataSink-->>PhysicalPlan: Iterator[WriteResult[FragmentMetadata]]
    end
    
    PhysicalPlan-->>WriteSink: All write results collected
    WriteSink->>WriteSink: results = write_df.to_pydict()
    WriteSink->>LanceDataSink: sink.finalize(results["write_results"])
    
    Note over LanceDataSink: finalize():<br/>- Extract all fragments from write_results<br/>- Create LanceOperation (Overwrite/Append)<br/>- Commit operation to create final dataset<br/>- Return stats as MicroPartition
    
    loop For each write_result in write_results
        LanceDataSink->>LanceDataSink: chain fragments from write_result.result
    end
    
    alt mode == "create" or "overwrite"
        LanceDataSink->>Lance: LanceOperation.Overwrite(schema, all_fragments)
    else mode == "append"
        LanceDataSink->>Lance: LanceOperation.Append(all_fragments)
    end
    
    LanceDataSink->>Lance: LanceDataset.commit(uri, operation, read_version, storage_options)
    Lance-->>LanceDataSink: committed_dataset
    LanceDataSink->>Lance: dataset.stats.dataset_stats()
    Lance-->>LanceDataSink: stats_dict
    
    LanceDataSink->>LanceDataSink: MicroPartition.from_pydict({num_fragments, num_deleted_rows, num_small_files, version})
    LanceDataSink-->>WriteSink: final_micropartition
    
    WriteSink->>WriteSink: to_logical_plan_builder(micropartition)
    WriteSink-->>DataFrame: DataFrame with write results
    DataFrame-->>User: final result DataFrame

Loading

dataframe.py

At daft/dataframe/dataframe.py with def write_lance()

class DataFrame:
    def __init__(self, builder: LogicalPlanBuilder) -> None:
        """Constructs a DataFrame according to a given LogicalPlan.

        Users are expected instead to call the classmethods on DataFrame to create a DataFrame.

        Args:
            plan: LogicalPlan describing the steps required to arrive at this DataFrame
        """
        if not isinstance(builder, LogicalPlanBuilder):
            if isinstance(builder, dict):
                raise ValueError(
                    "DataFrames should be constructed with a dictionary of columns using `daft.from_pydict`"
                )
            if isinstance(builder, list):
                raise ValueError(
                    "DataFrames should be constructed with a list of dictionaries using `daft.from_pylist`"
                )
            raise ValueError(f"Expected DataFrame to be constructed with a LogicalPlanBuilder, received: {builder}")

        self.__builder = builder
        self._result_cache: Optional[PartitionCacheEntry] = None
        self._preview = Preview(partition=None, total_rows=None)
        self._num_preview_rows = get_context().daft_execution_config.num_preview_rows

    @property
    def _builder(self) -> LogicalPlanBuilder:
        if self._result_cache is None:
            return self.__builder
        else:
            num_partitions = self._result_cache.num_partitions()
            size_bytes = self._result_cache.size_bytes()
            num_rows = self._result_cache.num_rows()

            # Partition set should always be set on cache entry.
            assert (
                num_partitions is not None and size_bytes is not None and num_rows is not None
            ), "Partition set should always be set on cache entry"

            return self.__builder.from_in_memory_scan(
                self._result_cache,
                self.__builder.schema(),
                num_partitions=num_partitions,
                size_bytes=size_bytes,
                num_rows=num_rows,
            )

    def _get_current_builder(self) -> LogicalPlanBuilder:
        """Returns the current logical plan builder, without any caching optimizations."""
        return self.__builder

    @property
    def _result(self) -> Optional[PartitionSet[PartitionT]]:
        if self._result_cache is None:
            return None
        else:
            return self._result_cache.value
    
    ...

@DataframePublicAPI
    def write_lance(
        self,
        uri: Union[str, pathlib.Path],
        mode: Literal["create", "append", "overwrite"] = "create",
        io_config: Optional[IOConfig] = None,
        **kwargs: Any,
    ) -> "DataFrame":
        from daft.dataframe.lance_data_sink import LanceDataSink

        sink = LanceDataSink(uri, self.schema(), mode, io_config, **kwargs)
        return self.write_sink(sink)

lance_data_sink.py

using daft/dataframe/lance_data_sink.py class LanceDataSink

class LanceDataSink(DataSink[list[lance.FragmentMetadata]]):
    ...
    def __init__(
            self,
            uri: str | pathlib.Path,
            schema: Schema,
            mode: Literal["create", "append", "overwrite"],
            io_config: IOConfig | None = None,
            **kwargs: Any,
        ) -> None:
            ...

            # Map Args/Kwargs to properties
            self._table_uri = str(uri)
            self._mode = mode
            self._io_config = get_context().daft_planning_config.default_io_config if io_config is None else io_config
            self._kwargs = kwargs # For 

            # Convert IO config back to storage_options
            self._storage_options = io_config_to_storage_options(self._io_config, self._table_uri)
            
            # Convert schema back to arrow 
            self._pyarrow_schema = pa.schema((f.name, f.dtype.to_arrow_dtype()) for f in schema)

            # Read Existing lance.LanceDataset from table uri
            try:
                table = lance.dataset(self._table_uri, storage_options=self._storage_options)
            except ValueError:
                table = None

            
            # Check that Dataset Schema and daft-converted pyarrow schema match. 
            self._version = 0 
            if table:
                table_schema = table.schema
                self._version = table.latest_version

                if self._pyarrow_schema != table_schema and not (self._mode == "overwrite"):
                    raise ValueError(
                        "Schema of data does not match table schema\n"
                        f"Data schema:\n{self._pyarrow_schema}\nTable Schema:\n{table_schema}"
                    )

            # Set internal daft schema
            self._schema = Schema._from_field_name_and_types(
                [
                    ("num_fragments", DataType.int64()),
                    ("num_deleted_rows", DataType.int64()),
                    ("num_small_files", DataType.int64()),
                    ("version", DataType.int64()),
                ]
            )

    # Daft Schema
    def schema(self) -> Schema:
        return self._schema

    # 
    def write(self, micropartitions: Iterator[MicroPartition]) -> Iterator[WriteResult[list[lance.FragmentMetadata]]]:
        """Writes fragments from the given micropartitions."""
        lance = self._import_lance()

        for micropartition in micropartitions:
            arrow_table = micropartition.to_arrow()
            bytes_written = arrow_table.nbytes
            rows_written = arrow_table.num_rows

            fragments = lance.fragment.write_fragments(
                arrow_table,
                dataset_uri=self._table_uri,
                mode=self._mode,
                storage_options=self._storage_options,
                **self._kwargs,
            )
            yield WriteResult(
                result=fragments,
                bytes_written=bytes_written,
                rows_written=rows_written,
            )

    def finalize(self, write_results: list[WriteResult[list[lance.FragmentMetadata]]]) -> MicroPartition:
        """Commits the fragments to the Lance dataset. Returns a DataFrame with the stats of the dataset."""
        from daft.dependencies import pa

        lance = self._import_lance()

        fragments = list(chain.from_iterable(write_result.result for write_result in write_results))

        if self._mode == "create" or self._mode == "overwrite":
            operation = lance.LanceOperation.Overwrite(self._pyarrow_schema, fragments)
        elif self._mode == "append":
            operation = lance.LanceOperation.Append(fragments)

        dataset = lance.LanceDataset.commit(
            self._table_uri, operation, read_version=self._version, storage_options=self._storage_options
        )
        stats = dataset.stats.dataset_stats()

        tbl = MicroPartition.from_pydict(
            {
                "num_fragments": pa.array([stats["num_fragments"]], type=pa.int64()),
                "num_deleted_rows": pa.array([stats["num_deleted_rows"]], type=pa.int64()),
                "num_small_files": pa.array([stats["num_small_files"]], type=pa.int64()),
                "version": pa.array([dataset.version], type=pa.int64()),
            }
        )
        return tbl

write_sink() in dataframe.py

whose methods are executed with DataFrame.write_sink()

@DataframePublicAPI
def write_sink(self, sink: "DataSink[WriteResultType]") -> "DataFrame":
    
    sink.start() # pass through in this case

    builder = self._builder.write_datasink(sink.name(), sink) # Routes through LogicalPlanBuilder 
    # builder is the result of LogicalPlanBuilder.write_lance()
    write_df = DataFrame(builder)
    write_df.collect()

    results = write_df.to_pydict()
    assert "write_results" in results
    micropartition = sink.finalize(results["write_results"])
    if micropartition.schema() != sink.schema():
        raise ValueError(
            f"Schema mismatch between the data sink's schema and the result's schema:\nSink schema:\n{sink.schema()}\nResult schema:\n{micropartition.schema()}"
        )
    # TODO(desmond): Connect the old and new logical plan builders so that a .explain() shows the
    # plan from the source all the way to the sink to the sink's results. In theory we can do this
    # for all other sinks too.
    write_plan_builder = to_logical_plan_builder(micropartition)
    return DataFrame(write_plan_builder)

write_lance in daft/logical/builder.py

which routes the sink operation through LogicalPlanBuilder.write_lance() from daft/logical/builder.py

class LogicalPlanBuilder:
    ...
    def write_lance(
        self,
        path: str | pathlib.Path,
        mode: str,
        io_config: IOConfig,
        kwargs: dict[str, Any] | None,
    ) -> LogicalPlanBuilder:
        columns_name = self.schema().column_names()
        builder = self._builder.lance_write(
            str(path),
            columns_name,
            mode,
            io_config,
            kwargs,
        )
        return LogicalPlanBuilder(builder)

    def write_datasink(self, name: str, sink: DataSink[Any]) -> LogicalPlanBuilder:
        builder = self._builder.datasink_write(name, sink)
        return LogicalPlanBuilder(builder)
    

lance_write in daft/execution/physical_plan.py

which runs def lance_write() from daft/execution/physical_plan.py

def lance_write(
    child_plan: InProgressPhysicalPlan[PartitionT],
    base_path: str,
    mode: str,
    io_config: IOConfig | None,
    kwargs: dict[str, Any] | None,
) -> InProgressPhysicalPlan[PartitionT]:
    """Write the results of `child_plan` into lance data files described by `write_info`."""
    yield from (
        step.add_instruction(
            execution_step.WriteLance(
                base_path=base_path,
                mode=mode,
                io_config=io_config,
                kwargs=kwargs,
            ),
        )
        if isinstance(step, PartitionTaskBuilder)
        else step
        for step in child_plan
    )

WriteLance in daft/execution/execution_step.py

which executes WriteLance in daft/execution/execution_step.py

@dataclass(frozen=True)
class WriteLance(SingleOutputInstruction):
    base_path: str
    mode: str
    io_config: IOConfig | None
    kwargs: dict[str, Any] | None

    def run(self, inputs: list[MicroPartition]) -> list[MicroPartition]:
        return self._write_lance(inputs)

    def _write_lance(self, inputs: list[MicroPartition]) -> list[MicroPartition]:
        [input] = inputs
        partition = self._handle_file_write(
            input=input,
        )
        return [partition]

    def run_partial_metadata(self, input_metadatas: list[PartialPartitionMetadata]) -> list[PartialPartitionMetadata]:
        assert len(input_metadatas) == 1
        return [
            PartialPartitionMetadata(
                num_rows=None,  # we can write more than 1 file per partition
                size_bytes=None,
            )
        ]

    def _handle_file_write(self, input: MicroPartition) -> MicroPartition:
        return recordbatch_io.write_lance(
            input,
            base_path=self.base_path,
            mode=self.mode,
            io_config=self.io_config,
            kwargs=self.kwargs,
        )

write_lance in daft/execution/record_batch_io.py

which FINALLY executes write_lance() in daft/execution/recordbatch_io.py

def write_lance(
    mp: MicroPartition,
    base_path: str,
    mode: str,
    io_config: IOConfig | None,
    kwargs: dict[str, Any] | None,
) -> MicroPartition:
    import lance

    from daft.io.object_store_options import io_config_to_storage_options

    io_config = get_context().daft_planning_config.default_io_config if io_config is None else io_config
    storage_options = io_config_to_storage_options(io_config, base_path)

    arrow_table = mp.to_arrow()

    fragments = lance.fragment.write_fragments(arrow_table, base_path, mode, storage_options=storage_options, **kwargs)

    mp = MicroPartition.from_pydict({"fragments": fragments})

    return mp

[everettVT]

I think this would be a great piece of source material as a prompt for a deep research prompted blog post.

[Jay-ju]

Good job, @everettVT. I noticed you've done some work on the daft lance connector. I'd like to know your thoughts. Let's see if we can cooperate to break down this work?