List column information of the dataset #

import daft
import os
import lance
from daft.io.object_store_options import io_config_to_storage_options
from daft.io import IOConfig
from daft.las.io import TOSConfig

# Set access key, secret key, and endpoint information
# Please replace with your access key and secret key

ak = "<your_ak>"  # Please replace with your access key
sk = "<your_sk>"  # Please replace with your secret key
endpoint = "https://tos-cn-beijing.ivolces.com"
region = "cn-beijing"

# Set key information as environment variables
os.environ["LAS_TOS_ACCESS_KEY"] = ak
os.environ["LAS_TOS_SECRET_KEY"] = sk
os.environ["TOS_ENDPOINT"] = endpoint

def list_dataset_columns(dataset_name=None, lance_tos_dir=None):
    """
    列出数据集的列名
    
    参数:
        dataset_name: LAS数据集名称
        lance_tos_dir: Lance数据集TOS路径
    """
    if dataset_name:
        # Read by dataset name
        print(f"\n读取LAS数据集: {dataset_name}")
        df = daft.read_las_dataset(name=dataset_name)
        print("\n数据集列名:")
        print(df.column_names)
        
        # Display a sample of the dataset
        print("\n数据集样例:")
        df.show(5)
    
    if lance_tos_dir:
        # Read directly via Lance path
        print(f"\n读取Lance数据集: {lance_tos_dir}")
        
        # Configure IO parameters using TOS configuration obtained from environment variables
        io_config = IOConfig(s3=TOSConfig.from_env().to_s3_config())
        
        # Generate Lance read configuration and build Lance client
        storage_options = io_config_to_storage_options(io_config, lance_tos_dir.replace("tos://", "s3://"))
        lance_ds = lance.dataset(uri=lance_tos_dir.replace("tos://", "s3://"), storage_options=storage_options)
        
        # Obtain schema of the Lance dataset
        schema = lance_ds.schema
        print("\nLance数据集schema:")
        print(schema)
        
        # List all column names
        print("\nLance数据集列名:")
        column_names = [field.name for field in schema]
        print(column_names)

def main():
    # Example: List column names of the LAS dataset
    # Replace with your dataset name
    las_dataset_name = "las_dataset_20251011_image_lance"
    list_dataset_columns(dataset_name=las_dataset_name)
    
    # Example: List column names of the Lance dataset
    # Replace with your TOS bucket name and path
    lance_tos_dir = "tos://xuxiaoliang-sh-test/lance/las_dataset_20251011_1_image.lance"
    list_dataset_columns(lance_tos_dir=lance_tos_dir)

if __name__ == "__main__":
    main()

Delete columns that are not needed #

##Please run source env.sh first to obtain relevant variable values
import os
import lance
from daft.io import IOConfig
from daft.io.object_store_options import io_config_to_storage_options
from daft.las.io import TOSConfig
from daft.las.infra.las_dataset import LasDatasetClient, LasDatasetConfig

# -------------------------- 1. Import core dependencies and configure environment variables (required for column deletion) --------------------------
# Obtain key configuration from environment variables (must be set in advance via environment variables or script)
TOS_LANCE_DIR = os.getenv("TOS_LANCE_DIR")  # Path of the Lance dataset in TOS (for example: tos://bucket/lance/dataset.lance)
DATASET_NAME = os.getenv("DATASET_NAME")              # LAS dataset name (used to refresh metadata)

# Load TOS access key (read from environment variables to avoid hardcoding)
os.environ["LAS_TOS_ACCESS_KEY"] = os.getenv("LAS_TOS_ACCESS_KEY")
os.environ["LAS_TOS_SECRET_KEY"] = os.getenv("LAS_TOS_SECRET_KEY")
os.environ["TOS_ENDPOINT"] = os.getenv("LAS_TOS_ENDPOINT")

# Initialize IO configuration (required for connecting to TOS, used for Lance dataset read/write)
io_config = IOConfig(s3=TOSConfig.from_env().to_s3_config())


# -------------------------- 2. Key features: Delete specified columns from the Lance dataset and refresh metadata --------------------------
def delete_embedding_columns():
    # 1. Process Lance path format (convert tos:// to s3:// to adapt to Lance SDK)
    if not TOS_LANCE_DIR:
        raise ValueError("请先设置环境变量 TOS_LANCE_DIR（Lance数据集在TOS的路径）")
    lance_path = TOS_LANCE_DIR.replace("tos://", "s3://")

    # 2. Connect to the Lance dataset (requires TOS storage configuration)
    storage_options = io_config_to_storage_options(io_config, lance_path)
    lance_ds = lance.dataset(uri=lance_path, storage_options=storage_options)
    print(f"成功连接Lance数据集：{lance_path}")
    print(f"当前数据集列：{lance_ds.schema.names}")

    # 3. Core operation: Delete specified columns (retain column names to be deleted in original code)
    columns_to_delete = ["num_rows", "image_embedding"]  # List of column names to be deleted
    # Filter columns that actually exist in the dataset (to avoid errors when deleting non-existent columns)
    existing_columns = [col for col in columns_to_delete if col in lance_ds.schema.names]

    if existing_columns:
        lance_ds.drop_columns(existing_columns)  # Perform column deletion operation
        print(f"已成功删除列：{existing_columns}")
    else:
        print(f"无需删除：指定的列 {columns_to_delete} 均不在数据集中")

    # 4. Refresh LAS console metadata (ensure the latest column structure is visible in the control panel after deletion)
    if DATASET_NAME:
        client = LasDatasetClient(config=LasDatasetConfig.from_io_config(io_config))
        body = {"DatasetName": DATASET_NAME}
        client.api_client.call_api(
            method="POST",
            params={},
            headers={},
            action="RefreshDatasetMetadata",
            body=body,
        )
        print(f"已刷新LAS数据集 {DATASET_NAME} 的元数据")
    else:
        print("未设置 DATASET_NAME，跳过元数据刷新步骤")


# -------------------------- 3. Entry point --------------------------
if __name__ == "__main__":
    delete_embedding_columns()