|
2 | 2 |
|
3 | 3 | Create Lance tables from DataFrames using the DataSource V2 API. |
4 | 4 |
|
| 5 | +## Basic DataFrame Creation |
| 6 | + |
5 | 7 | === "Python" |
6 | 8 | ```python |
7 | 9 | # Create DataFrame |
@@ -55,4 +57,159 @@ Create Lance tables from DataFrames using the DataSource V2 API. |
55 | 57 |
|
56 | 58 | // Write as new table using catalog |
57 | 59 | df.writeTo("users").create(); |
58 | | - ``` |
| 60 | + ``` |
| 61 | + |
| 62 | +## Creating Tables with Vector Columns |
| 63 | + |
| 64 | +Lance supports vector (embedding) columns for AI workloads. These columns are stored internally as Arrow `FixedSizeList[n]` where `n` is the vector dimension. Since Spark DataFrames don't have a native fixed-size array type, you need to add metadata to your schema fields to indicate that an `ArrayType(FloatType)` or `ArrayType(DoubleType)` should be converted to Arrow FixedSizeList. |
| 65 | + |
| 66 | +The metadata key `"arrow.fixed-size-list.size"` with a value like `128` tells the Lance-Spark connector to convert that array column to a `FixedSizeList[128]` during write operations. |
| 67 | + |
| 68 | +### Supported Types |
| 69 | + |
| 70 | +- **Element Types**: `FloatType` (float32), `DoubleType` (float64) |
| 71 | +- **Array Requirements**: |
| 72 | + - Must have `containsNull = false` |
| 73 | + - Column must be non-nullable |
| 74 | + - All arrays must have exactly the specified dimension |
| 75 | + |
| 76 | +### Examples |
| 77 | + |
| 78 | +=== "Python" |
| 79 | + ```python |
| 80 | + from pyspark.sql.types import StructType, StructField, IntegerType, ArrayType, FloatType |
| 81 | + from pyspark.sql.types import Metadata |
| 82 | + |
| 83 | + # Create schema with vector column |
| 84 | + vector_metadata = {"arrow.fixed-size-list.size": 128} |
| 85 | + schema = StructType([ |
| 86 | + StructField("id", IntegerType(), False), |
| 87 | + StructField("embeddings", ArrayType(FloatType(), False), False, vector_metadata) |
| 88 | + ]) |
| 89 | + |
| 90 | + # Create DataFrame with vector data |
| 91 | + import numpy as np |
| 92 | + data = [(i, np.random.rand(128).astype(np.float32).tolist()) for i in range(100)] |
| 93 | + df = spark.createDataFrame(data, schema) |
| 94 | + |
| 95 | + # Write to Lance format |
| 96 | + df.writeTo("vectors_table").create() |
| 97 | + ``` |
| 98 | + |
| 99 | +=== "Scala" |
| 100 | + ```scala |
| 101 | + import org.apache.spark.sql.types._ |
| 102 | + |
| 103 | + // Create metadata for vector column |
| 104 | + val vectorMetadata = new MetadataBuilder() |
| 105 | + .putLong("arrow.fixed-size-list.size", 128) |
| 106 | + .build() |
| 107 | + |
| 108 | + // Create schema with vector column |
| 109 | + val schema = StructType(Array( |
| 110 | + StructField("id", IntegerType, false), |
| 111 | + StructField("embeddings", ArrayType(FloatType, false), false, vectorMetadata) |
| 112 | + )) |
| 113 | + |
| 114 | + // Create DataFrame with vector data |
| 115 | + import scala.util.Random |
| 116 | + val data = (0 until 100).map { i => |
| 117 | + (i, Array.fill(128)(Random.nextFloat())) |
| 118 | + } |
| 119 | + val df = spark.createDataFrame(data).toDF("id", "embeddings") |
| 120 | + |
| 121 | + // Write to Lance format |
| 122 | + df.writeTo("vectors_table").create() |
| 123 | + ``` |
| 124 | + |
| 125 | +=== "Java" |
| 126 | + ```java |
| 127 | + import org.apache.spark.sql.types.*; |
| 128 | + |
| 129 | + // Create metadata for vector column |
| 130 | + Metadata vectorMetadata = new MetadataBuilder() |
| 131 | + .putLong("arrow.fixed-size-list.size", 128) |
| 132 | + .build(); |
| 133 | + |
| 134 | + // Create schema with vector column |
| 135 | + StructType schema = new StructType(new StructField[] { |
| 136 | + DataTypes.createStructField("id", DataTypes.IntegerType, false), |
| 137 | + DataTypes.createStructField("embeddings", |
| 138 | + DataTypes.createArrayType(DataTypes.FloatType, false), |
| 139 | + false, vectorMetadata) |
| 140 | + }); |
| 141 | + |
| 142 | + // Create DataFrame with vector data |
| 143 | + List<Row> rows = new ArrayList<>(); |
| 144 | + Random random = new Random(); |
| 145 | + for (int i = 0; i < 100; i++) { |
| 146 | + float[] vector = new float[128]; |
| 147 | + for (int j = 0; j < 128; j++) { |
| 148 | + vector[j] = random.nextFloat(); |
| 149 | + } |
| 150 | + rows.add(RowFactory.create(i, vector)); |
| 151 | + } |
| 152 | + Dataset<Row> df = spark.createDataFrame(rows, schema); |
| 153 | + |
| 154 | + // Write to Lance format |
| 155 | + df.writeTo("vectors_table").create(); |
| 156 | + ``` |
| 157 | + |
| 158 | +### Creating Multiple Vector Columns |
| 159 | + |
| 160 | +You can create DataFrames with multiple vector columns, each with different dimensions: |
| 161 | + |
| 162 | +=== "Python" |
| 163 | + ```python |
| 164 | + from pyspark.sql.types import DoubleType |
| 165 | + |
| 166 | + # Create schema with multiple vector columns |
| 167 | + text_metadata = {"arrow.fixed-size-list.size": 384} |
| 168 | + image_metadata = {"arrow.fixed-size-list.size": 512} |
| 169 | + |
| 170 | + schema = StructType([ |
| 171 | + StructField("id", IntegerType(), False), |
| 172 | + StructField("text_embeddings", ArrayType(FloatType(), False), False, text_metadata), |
| 173 | + StructField("image_embeddings", ArrayType(DoubleType(), False), False, image_metadata) |
| 174 | + ]) |
| 175 | + |
| 176 | + # Create DataFrame with multiple vector columns |
| 177 | + data = [ |
| 178 | + (i, |
| 179 | + np.random.rand(384).astype(np.float32).tolist(), |
| 180 | + np.random.rand(512).tolist()) |
| 181 | + for i in range(100) |
| 182 | + ] |
| 183 | + df = spark.createDataFrame(data, schema) |
| 184 | + |
| 185 | + # Write to Lance format |
| 186 | + df.writeTo("multi_vectors_table").create() |
| 187 | + ``` |
| 188 | + |
| 189 | +### Vector Indexing |
| 190 | + |
| 191 | +After creating tables with vector columns, you can create vector indexes for efficient similarity search using the Lance Python API: |
| 192 | + |
| 193 | +```python |
| 194 | +import lance |
| 195 | +import numpy as np |
| 196 | + |
| 197 | +# Open the dataset |
| 198 | +ds = lance.dataset("/path/to/vectors_table.lance") |
| 199 | + |
| 200 | +# Create a vector index on the embeddings column |
| 201 | +ds.create_index( |
| 202 | + "embeddings", |
| 203 | + index_type="IVF_PQ", |
| 204 | + num_partitions=256, |
| 205 | + num_sub_vectors=16 |
| 206 | +) |
| 207 | + |
| 208 | +# Perform similarity search |
| 209 | +query_vector = np.random.rand(128).astype(np.float32) |
| 210 | +results = ds.to_table( |
| 211 | + nearest={"column": "embeddings", "q": query_vector, "k": 10} |
| 212 | +).to_pandas() |
| 213 | +``` |
| 214 | + |
| 215 | +Note: When reading vector columns back into Spark DataFrames, they are automatically converted to regular `ArrayType(FloatType)` or `ArrayType(DoubleType)` for compatibility with Spark operations. |
0 commit comments