Function read_dataset_from_parquet

tests/test-dataset.py:433–484

Read dataset from parquet files with a strict limit to prevent memory issues.

(dataset_name, max_vectors=50000)

def read_dataset_from_parquet(dataset_name, max_vectors=50000):
    """Read dataset from parquet files with a strict limit to prevent memory issues."""
    metadata = datasets[dataset_name]

    datasets_dir = "datasets"
    dataset_dir = os.path.join(datasets_dir, dataset_name)

    path = os.path.join(dataset_dir, "test0.parquet")
    if not os.path.exists(path):
        raise FileNotFoundError(
            f"No parquet files found for dataset {dataset_name}. "
            f"Please place your parquet files in: {os.path.abspath(dataset_dir)}"
        )

    vectors = []
    file_count = 0
    
    print(f"Reading dataset with limit of {max_vectors} vectors to prevent memory issues...")
    
    while os.path.exists(path) and len(vectors) < max_vectors:
        # Read the full parquet file but limit processing
        df = pd.read_parquet(path)
        
        # Limit rows to prevent memory overload
        rows_to_read = min(len(df), max_vectors - len(vectors))
        df = df.head(rows_to_read)
        
        emb_col = "dense_values" if "dense_values" in df.columns else metadata["embeddings"]
        
        if metadata["id"] is not None:
            id_col = metadata["id"]
            dataset_chunk = df[[id_col, emb_col]].values.tolist()
        else:
            dataset_chunk = list(enumerate(df[emb_col].values))

        for row in dataset_chunk:
            if len(vectors) >= max_vectors:
                break
            vector = pre_process_vector(row[0], row[1])
            vectors.append(vector)

        file_count += 1
        path = os.path.join(dataset_dir, f"test{file_count}.parquet")
        
        # Clear dataframe from memory
        del df

    if not vectors:
        raise ValueError(f"No data found in dataset {dataset_name}")

    print(f"Loaded {len(vectors)} vectors from dataset {dataset_name}")
    return vectors

def read_single_parquet_file(path, dataset_name, file_index, base_id, quick_test=False):
    """Read and process a single parquet file"""