update lance

Author: allisonwang-db

pyspark_datasources/lance.py

@@ -2,7 +2,7 @@
 import pyarrow as pa
 
 from dataclasses import dataclass
-from typing import Iterator
+from typing import Iterator, List
 from pyspark.sql.datasource import DataSource, DataSourceArrowWriter, WriterCommitMessage
 from pyspark.sql.pandas.types import to_arrow_schema
 
@@ -22,7 +22,7 @@ class LanceSink(DataSource):
 
     Create a Spark dataframe with 2 partitions:
 
-    >>> df = spark.range(0, 3, 1, 2)
+    >>> df = spark.createDataFrame([(1, "a"), (2, "b"), (3, "c")], schema="id int, value string")
 
     Save the dataframe in lance format:
 
@@ -58,7 +58,7 @@ def writer(self, schema, overwrite: bool):
 
 @dataclass
 class LanceCommitMessage(WriterCommitMessage):
-    fragment: lance.FragmentMetadata
+    fragments: List[lance.FragmentMetadata]
 
 
 class LanceWriter(DataSourceArrowWriter):
@@ -78,18 +78,12 @@ def _get_read_version(self):
             return None
 
     def write(self, iterator: Iterator[pa.RecordBatch]):
-        from pyspark import TaskContext
-
-        context = TaskContext.get()
-        assert context is not None, "Unable to get TaskContext"
-        task_id = context.taskAttemptId()
-
         reader = pa.RecordBatchReader.from_batches(self.arrow_schema, iterator)
-        fragment = lance.LanceFragment.create(self.uri, reader, fragment_id=task_id, schema=self.arrow_schema)
-        return LanceCommitMessage(fragment=fragment)
+        fragments = lance.fragment.write_fragments(reader, self.uri, schema=self.arrow_schema)
+        return LanceCommitMessage(fragments=fragments)
 
     def commit(self, messages):
-        fragments = [msg.fragment for msg in messages]
+        fragments = [fragment for msg in messages for fragment in msg.fragments ]
         if self.read_version:
             # This means the dataset already exists.
             op = lance.LanceOperation.Append(fragments)