基于Notebook+PySpark+Lance实现高效图文混存

登录云原生数据仓库AnalyticDB MySQL控制台，在左上角选择集群所在地域。在左侧导航栏，单击集群列表，然后单击目标集群ID。

单击作业开发 > Notebook开发。确保已完成如下准备工作，然后单击进入DMS Notebook。

新建工作空间。

引入数据源。

创建Spark集群。

创建并启动Notebook会话。

单击按钮，然后单击+ > 新建Notebook文件。

将Cell的语言类型设置为Python，执行以下代码，下载Python依赖。

!pip install pyarrow==19.0.1
!pip install pylance==0.23.2
!pip install oss2==2.19.1
!pip install pandas==2.2.3
!pip install torch==2.7.0
!pip install torchvision==0.22.0
!pip install pillow==11.2.1

将Cell的语言类型设置为Python，替换以下代码中的相关配置参数并执行，将图片数据导入Lance。

import os
import oss2
import pandas as pd
import pyarrow as pa
import lance

# Bucket名称
bucket_name = 'testBucketName'
# Bucket的Endpoint
endpoint = 'oss-cn-hangzhou-internal.aliyuncs.com'
# 阿里云账号或者具备OSS访问权限的RAM用户的AccessKey ID和AccessKey Secret
auth = oss2.Auth('AK', 'SK')
bucket = oss2.Bucket(auth, endpoint, bucket_name)

# 图片数据存放的目录前缀
prefix = 'lanceData/'

storage_options = {
    # Bucket所属地域
    "region": "cn-hangzhou",
    # Bucket的Endpoint
    "endpoint": "https://testBucketName.oss-cn-hangzhou-internal.aliyuncs.com",
    # 阿里云账号或者具备OSS访问权限的RAM用户的AccessKey ID和AccessKey Secret
    "access_key_id": "ak",
    "secret_access_key": "sk",
    "virtual_hosted_style_request": "True"
}

data = []
for obj in oss2.ObjectIterator(bucket, prefix=prefix):
    if obj.key.endswith(('.png', '.jpg', '.jpeg')):
        image_id = os.path.splitext(os.path.basename(obj.key))[0]
        img_binary = bucket.get_object(obj.key).read()
        data.append({'image_id': image_id, 'image_data': img_binary})

df = pd.DataFrame(data)

schema = pa.schema([pa.field("image_id", pa.string()), pa.field("image_data", pa.binary())])
table = pa.Table.from_pandas(df, schema=schema)
# 图片Lance数据存放的OSS路径，协议需为s3
uri = "s3://testBucketName/lance_data/lance_image_dataset.lance"
lance.write_dataset(table, uri, storage_options=storage_options)

ds = lance.dataset(uri, storage_options=storage_options)

将Cell的语言类型设置为Python，替换以下代码中的相关配置参数并执行，将文本数据导入Lance。

import pandas as pd
import pyarrow as pa
import lance
import fsspec
import json
import oss2
import json

storage_options = {
    # Bucket所属地域
    "region": "cn-hangzhou",
    # Bucket的Endpoint
    "endpoint": "https://testBucketName.oss-cn-hangzhou-internal.aliyuncs.com",
    # 阿里云账号或者具备OSS访问权限的RAM用户的AccessKey ID和AccessKey Secret
    "access_key_id": "AK",
    "secret_access_key": "SK",
    "virtual_hosted_style_request": "True"
}


# 初始化认证信息
auth = oss2.Auth('AK', 'SK')
bucket = oss2.Bucket(
    auth,
    'https://oss-cn-hangzhou-internal.aliyuncs.com', 
    'testBucketName' # Bucket名称
)

# 取文件内容
obj = bucket.get_object('Lance/data/testdata.json')
text_data = json.load(obj)
# 2. 构造DataFrame并后续数据处理
schema = pa.schema([
    pa.field("docid", pa.string()),
    pa.field("text", pa.string()),
    pa.field("image_id1", pa.string()),
    pa.field("image_id2", pa.string()),
    pa.field("image_id3", pa.string())
])

df = pd.DataFrame(text_data)
table = pa.Table.from_pandas(df, schema=schema)
# 文本Lance数据存放的OSS路径，协议需为s3
uri = "s3://testBucketName/Lance/lance_text_dataset.lance"

lance.write_dataset(table, uri, storage_options=storage_options)

ds = lance.dataset(uri, storage_options=storage_options)
print(ds.to_table().to_pandas())

将Cell的语言类型设置为Python，替换以下代码中的相关配置参数并执行，将图片与文本数据关联整合，存储为同时包含图片与文本的综合数据集。

from pyspark.sql.dataframe import DataFrame

from pyspark.sql import SparkSession, DataFrame
from pyspark.sql import functions as F
from pyspark.sql.types import StringType
import lance
import json

storage_options = {
    # Bucket所属地域
    "region": "cn-hangzhou",
    # Bucket的Endpoint
    "endpoint": "https://testBucketName.oss-cn-hangzhou-internal.aliyuncs.com",
    # 阿里云账号或者具备OSS访问权限的RAM用户的AccessKey ID和AccessKey Secret
    "access_key_id": "ak",
    "secret_access_key": "sk",
    "virtual_hosted_style_request": "True"
}
class VisionDataProcessor:
    def __init__(self, app_name="Lance Vision Data Join and Format Example"):
        self.spark = SparkSession.builder.config("spark.driver.memory", "32g") \
        .config("spark.sql.catalog.lance","com.lancedb.lance.spark.LanceCatalog") \
        .config("spark.sql.catalog.lance.aws_region","cn-hangzhou") \
         # Bucket的Endpoint
        .config("spark.sql.catalog.lance.aws_endpoint","https://testBucketName.oss-cn-hangzhou-internal.aliyuncs.com") \
        # 阿里云账号或者具备OSS访问权限的RAM用户的AccessKey ID和AccessKey Secret
        .config("spark.sql.catalog.lance.access_key_id","ak") \
        .config("spark.sql.catalog.lance.secret_access_key","sk") \
        .config("spark.sql.catalog.lance.virtual_hosted_style_request","True") \
        .config("spark.driver.cores", "8").config("spark.rpc.message.maxSize","1024").appName(app_name).getOrCreate()

    def get_lance_dataset(self, uri) -> DataFrame:
        # 读取 Lance 数据集
        data = spark.read \
            .format("lance") \
            .load(uri)
        return data

    def join_vision_data(self) -> DataFrame:
        # 图片Lance数据和文本Lance数据存放的OSS路径，协议需为s3
        df_text = self.get_lance_dataset(uri='s3://testBucketName/lance_data/lance_text_dataset.lance')
        df_image = self.get_lance_dataset('s3://testBucketName/lance_data/lance_image_dataset.lance')

        df_text = df_text.join(
            df_image.select(F.col("image_id").alias("image_id1"), F.col("image_data")),
            on="image_id1",
            how="left"
        ).withColumnRenamed("image_data", "image1_byte")

        df_text = df_text.join(
            df_image.select(F.col("image_id").alias("image_id2"), F.col("image_data")),
            on="image_id2",
            how="left"
        ).withColumnRenamed("image_data", "image2_byte")

        df_text = df_text.join(
            df_image.select(F.col("image_id").alias("image_id3"), F.col("image_data")),
            on="image_id3",
            how="left"
        ).withColumnRenamed("image_data", "image3_byte")

        vision_df = df_text.select(
            "docid",
            "text",
            F.col("image1_byte").alias("image1"),
            F.col("image2_byte").alias("image2"),
            F.col("image3_byte").alias("image3")
        )
        return vision_df

    def format_to_training_data(self, df: DataFrame) -> DataFrame:
        def format_prompt(text, image1, image2, image3):
            images = [image.hex() for image in (image1, image2, image3) if image]
            alpaca_format = {"instruction": text, "input": images, "output": ""}
            return json.dumps(alpaca_format)

        format_prompt_udf = F.udf(format_prompt, StringType())
        result_str = df.withColumn("formatted_prompt", format_prompt_udf("text", "image1", "image2", "image3"))
        return result_str

    def process_and_save(self, json_output_path="training_data.json"):
        vision_df = self.join_vision_data()
        formatted_df = self.format_to_training_data(vision_df)

        formatted_df.show(truncate=True)

        # 保存为JSON文件
        formatted_prompts = [row.formatted_prompt for row in formatted_df.collect()]
        with open(json_output_path, "w") as f:
            json.dump(formatted_prompts, f, indent=2)

将Cell的语言类型设置为Python，执行以下代码，构造训练数据集。

from torch.utils.data import Dataset, DataLoader
from PIL import Image
import torch
from torchvision import transforms
import io
import json 

class VLMDataset(Dataset):
    """
    Custom dataset for VLM training.
    """

    def __init__(self, df, transform=None):
        """
        Initialize the dataset.
        """
        self.rows = df.select("formatted_prompt").collect()
        self.transform = transform or transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor()
        ])

    def __len__(self):
        """
        Get the total number of samples in the dataset.
        """
        return len(self.rows)

    def __getitem__(self, idx):
        """
        Get a sample from the dataset.
        """
        row = json.loads(self.rows[idx].formatted_prompt)
        text = row["instruction"]
        images_data = row["input"]
        images = []

        for image_hex in images_data:
            img_data = bytes.fromhex(image_hex)
            try:
                image = Image.open(io.BytesIO(img_data))
                if self.transform:
                    image = self.transform(image)
                images.append(image)
            except Exception as e:
                print(f"Error processing image: {e}")
                images.append(torch.zeros(3, 224, 224))  # Placeholder for missing images

        while len(images) < 3:
            images.append(torch.zeros(3, 224, 224))

        return text, torch.stack(images)


if __name__ == "__main__":
    processor = VisionDataProcessor()
    formatted_df = processor.format_to_training_data(processor.join_vision_data())

    dataset = VLMDataset(formatted_df)
    dataloader = DataLoader(dataset, batch_size=8, shuffle=False)

    for batch in dataloader:
        texts, images = batch
        print('Texts:', texts)
        print('Images batch shape:', images.shape)

    processor.spark.stop()

执行结果如下：

构造了一个3*3*3*224*224（批次*图片数量*图片大小[RGB，长，宽]）的数据集。