When using PySpark to process Lance datasets (handled via pylance), the most common performance killer is involving large fields in Spark shuffle (data repartition exchange). For example, download_html, long text, binary data, vectors, and more. This article presents a proven engineering paradigm: only shuffle a few bytes of row address, push down the reading of large fields to within the partition, and use Lance's Fragment.take() to read by row address. This approach ensures correct data distribution while almost completely eliminating network and disk I/O for large fields.

Why large fields make shuffle expensive #

Spark's repartition, join, groupBy, and other wide dependency operators trigger shuffle, and the cost of shuffle is directly proportional to the volume of data being moved. When a row contains a large field of tens of KB or even MB:

• Each row must serialize this large field and write it to the shuffle file, causing disk I/O to surge.
• Network bandwidth is saturated during cross-Executor transfers.
• The Python (PySpark) side also bears additional serialization/deserialization and memory pressure, making it prone to OOM or frequent GC.

Why parquet/ORC cannot solve this problem #

Parquet/ORC is a format designed for "batch scanning + column pruning"; the smallest read unit is a row group, and there is no stable, cross-task transferable "row address", nor is there an interface like Fragment.take(offsets, columns=[...]) that enables constant-time random access.

Even if only the file path and row number are shuffled (data repartition exchange), downstream still needs to rescan the entire table to retrieve the target large field, and then perform a join to assemble the large field into the result table—this means at least two full I/O operations, making the cost even higher.

The cornerstone of Lance: row address and Fragment #

A Lance dataset consists of several Fragments (horizontal partitions), each Fragment has an auto-incrementing uint32 identifier, and is composed of one or more data files.
The row address _rowaddr is the physical location of a row in the table, represented as a 64-bit identifier formed by combining two 32-bit values:

Therefore, decoding _rowaddr to obtain the Fragment and offset only requires two bitwise operations:

fragment_id = rowaddr >> 32          # High 32 bits: confirm Fragment number
row_offset  = rowaddr & 0xFFFFFFFF   # Lower 32 bits: Confirm row offset within Fragment

This addressing scheme natively supports constant-time random access: once the row address is obtained, you can directly locate the Fragment and the offset within the file, and use an optimized columnar reader to fetch data in batches.

Lance's shuffle solution for large fields #

Lance fills two gaps in traditional data tables:

1. _rowaddr is a relatively stable, externally transferable physical row address.
2. Fragment.take() provides low-overhead random point lookup and column pruning.

In other words, shuffling only row addresses and fetching large fields on demand is not a Spark trick, but a capability unlocked by Lance's columnar format that natively supports row-level random access.

Why Lance is faster when sorting before fetching data #

This is the essence of the practical solution, derived from the physical meaning of _rowaddr:

1. Aggregation by Fragment: rows with the same upper 32 bits in _rowaddr belong to the same Fragment. Grouping them by Fragment allows a single Fragment object to be reused repeatedly, avoiding reopening files and reconstructing the reader for each row.
2. Contiguity within partitions: the lower 32 bits of _rowaddr are the incrementing offset within the Fragment. After sorting by _rowaddr, rows of the same Fragment are contiguously increasing within the partition, enabling effective hits on the Lance reader cache and OS page cache during reads, turning many random reads into near-sequential reads.

Core logic #

The essence of this paradigm is summed up in one sentence: Let Spark handle "address scheduling", let Lance handle "content reading"; let lightweight columns go through shuffle, and let heavy fields be read locally within partitions as needed.

1. On the Spark side, only _rowaddr/_fragid—these few-byte integers—are shuffled, completely decoupling shuffle volume from large fields.
2. repartition(_fragid) + sortWithinPartitions(_rowaddr) leverage the physical meaning of row addresses, enabling Fragment reuse and cache hits.
3. Fragment.take() reads large fields within partitions on demand, in batches, and in near-sequential order.

For Lance datasets containing large text, binary, vector, and other heavy fields, this approach typically reduces shuffle I/O by one to two orders of magnitude, making it a standard pattern worth solidifying in PySpark + Lance engineering practice.

Code implementation #

Anti-pattern vs recommended paradigm #

First, compare the two approaches to intuitively see the difference.

Anti-pattern illustration (do not write like this):

# ❌ Anti-pattern: download_html is brought into the DataFrame from the very beginning,
#    Any subsequent repartition / join / groupBy will move the large field around
big = ds_df.select("_rowaddr", "download_html", "label")
big = big.repartition("label")          # Large fields are involved in Shuffle, causing disk and network I/O to explode
result = big.join(other, on="key")      # Shuffle again, making things even worse

Step 1: Only read _rowaddr / _fragid, these lightweight columns #

We only read the "row address" into Spark; large fields are completely excluded.

import lance
import pyspark.sql.functions as F
from pyspark.sql import SparkSession
from pyspark.sql.types import StructType, StructField, LongType, StringType

# Object storage (TOS S3) connection configuration, used by Lance to read remote datasets
STORAGE_OPTIONS ={
    "aws_region": "cn-beijing",
    "aws_endpoint": "https://your-bucket.tos-s3-cn-beijing.ivolces.com",
    "access_key_id": "",
    "secret_access_key": "",
    "aws_virtual_hosted_style_request": "true"
}

# Lance dataset path (TOS S3 address)
DATASET_URI = "s3://your-bucket/your-dataset-path"

# Obtain the current dataset version number to ensure the entire task reads the same snapshot and avoid version drift
DATASET_VERSION = lance.dataset(DATASET_URI, storage_options=STORAGE_OPTIONS).version

# Initialize Spark session
spark = SparkSession.builder.appName("lance-avoid-bigfield-Shuffle").getOrCreate()

# Step 1: Read only lightweight addressing fields, do not load large fields
# _fragid: Fragment (horizontal shard) number, identifies the shard to which the row belongs
# _rowaddr: Row address, 64-bit integer, uniquely locates the physical position of any row in the table
# Purpose: In subsequent Shuffle (data repartition exchange), only row addresses (8 bytes per row) are moved,
addr_df = spark.read.format("lance").load(DATASET_URI).select("_fragid", "_rowaddr")

Step 2：repartition(_fragid) + sortWithinPartitions(_rowaddr) #

This step is key to ensuring locality:

# Step 2: Sort by physical address, partition by Fragment
addr_df = (
    addr_df
    .repartition("_fragid")            # All rows of the same Fragment are shuffled to the same partition
    .sortWithinPartitions("_rowaddr")  # Rows within a partition are arranged consecutively by physical address, facilitating reuse and cache hits
)
# To control parallelism, use .repartition(num_partitions, "_fragid")

• repartition("_fragid"): Ensures that all _rowaddr belonging to the same Fragment are placed in the same Partition, allowing downstream to reuse the same Fragment object.
• sortWithinPartitions("_rowaddr"): Since _rowaddr = (Fragment_id << 32) | offset, sorting by _rowaddr is equivalent to sorting first by Fragment, then by offset within the Fragment. As a result, data for the same Fragment is contiguous and increasing within the partition, turning random reads into near-sequential reads during access.

Step 3: Use Fragment.take() within each partition to batch fetch large fields #

With the guarantee of "aggregation by Fragment + contiguous sorting of row addresses within partitions", each Partition only needs to reuse a single Fragment object—after accumulating a batch of row offsets, call Fragment.take() to batch read multiple large fields at once, merging N individual I/O operations into one, significantly reducing remote storage access overhead.

# Step 3: Define the fetch_big_field function
# part_iter: The iterator for Partition rows passed in by Spark
# dataset_uri: The address of the Lance dataset in TOS S3
# dataset_version: The version number, locks the snapshot to prevent reading newly written data during processing
# columns: The names of the large field columns to retrieve; here, download_html is selected
# batch_size: The maximum number of rows to accumulate for unified retrieval each time; here, 512
def fetch_big_field(part_iter, dataset_uri, dataset_version,
                    columns=("download_html",), batch_size=512):
    
    # Each Python Worker process independently initializes the Dataset object (reads metadata). The Dataset must not be placed in the outer closure to avoid Spark serialization transmission issues
    ds = lance.dataset(dataset_uri, version=dataset_version) 
    columns = list(columns) # Convert the columns tuple to a list

    cur_frag_id = None  # Initialize the current Fragment ID being processed
    cur_Fragment = None # Initialize the current Fragment object
    buf_offsets = []    # Initialize the list of offsets for a batch
    buf_addrs = []      # Initialize the _rowaddr corresponding to buf_offsets, used for reconstruction

   # Define the flush subfunction, accumulate a batch of offsets and perform a single take, returning a list of (rowaddr, *cols)
    def flush():
        if not buf_offsets:
            return []
        # Key point: Retrieve column data corresponding to a batch of offsets to form tbl
        tbl = cur_Fragment.take(buf_offsets, columns=columns) 
         # Convert tbl to a regular Python list
        col_lists = [tbl.column(c).to_pylist() for c in columns]
       
        #Combine row addresses and column values into tuples by row to form a list of row objects
        rows = [(addr, *vals) for addr, *vals in zip(buf_addrs, *col_lists)]
        buf_offsets.clear() # Clear the buffer to prepare for the next batch
        buf_addrs.clear()
        return rows
    
    for row in part_iter:
        rowaddr = row["_rowaddr"]
        frag_id = row["_fragid"]
        # For each incoming row, first calculate its offset within the Fragment (lower 32 bits)
        offset  = rowaddr & 0xFFFFFFFF # Lower 32 bits (must be a 32-bit mask!)

        # If frag_id changes: flush the buffered batch of the previous Fragment before switching Fragment
        if frag_id != cur_frag_id:
            yield from flush()
            cur_frag_id = frag_id
            cur_Fragment = ds.get_Fragment(frag_id) # Fragment reuse: Only open the same Fragment once
        
        # Accumulate offsets and row addresses into the buffer
        buf_offsets.append(offset)
        buf_addrs.append(rowaddr)
        
        # Buffer reaches 512 entries: flush
        if len(buf_offsets) >= batch_size:
            yield from flush()    

    # Flush the final batch
    yield from flush()

Connect it to the sorted addr_df:

# Step 4: Call fetch_big_field for each Partition
result_rdd = addr_df.rdd.mapPartitions(
    lambda it: fetch_big_field(
        it, DATASET_URI, DATASET_VERSION,
        columns=("download_html",), batch_size=512,
    )
)

# Step 5: Convert rdd back to DataFrame for subsequent Spark SQL or structured operations
result_schema = StructType([
    StructField("_rowaddr", LongType(), False),
    StructField("download_html", StringType(), True),
])
result_df = spark.createDataFrame(result_rdd, schema=result_schema)

• _rowaddr is a physical address, not a stable ID: The row address will change with data reorganization such as compaction, update, delete. Be sure to fix version (as in DATASET_VERSION above) at the start of the job to ensure address validity throughout the job; if the dataset may change during the job, read from a fixed snapshot. For stable identification across versions, use Lance's stable row ID (_rowid, requires enabling enable_stable_row_ids at write time), but its access path differs from offset access via _rowaddr.
• **take The input parameter to take is the offset within Fragment: **The argument passed to Fragment.take() is row_offset (local offset within the Fragment), not the global row number, nor _rowaddr itself. Be sure to decode with & 0xFFFFFFFF first.
• Data skew: If some Fragment have far more matching rows than others, partition skew may occur after repartition("_fragid"). To mitigate this, you can apply salting or explicitly specify the number of partitions based on _fragid and _rowaddr, but ensure that rows from the same Fragment are still aggregated as much as possible to maintain reuse benefits.
• PySpark serialization: The lance.dataset referenced in the partition function must be created within the function body; do not pass pre-constructed Dataset/Fragment objects as closure variables to the Executor (not serializable). Only basic types such as uri and version should be passed.