Function generate_brute_force_results

tests/test-dataset.py:319–411

Generate brute force results by processing dataset in chunks

(dataset_name, quick_test=False)

def generate_brute_force_results(dataset_name, quick_test=False):
    """Generate brute force results by processing dataset in chunks"""
    print("Generating brute force results using chunk processing...")
    
    dataset_config = datasets[dataset_name]
    dataset_dir = os.path.join("datasets", dataset_name)
    
    # Determine sample size
    k = 10 if quick_test else 100
    reservoir = []
    total_vectors = 0
    file_count = 0
    
    # Pass 1: Reservoir sampling to select query vectors
    random.seed(42)
    while True:
        path = os.path.join(dataset_dir, f"test{file_count}.parquet")
        if not os.path.exists(path):
            break
            
        df = pd.read_parquet(path)
        for index, row in df.iterrows():
            total_vectors += 1
            # Get vector data
            id_val = row[dataset_config["id"]] if dataset_config["id"] else index
            embedding = row[dataset_config["embeddings"]]
            vector = pre_process_vector(id_val, embedding)
            
            # Reservoir sampling
            if len(reservoir) < k:
                reservoir.append((vector["id"], vector["dense_values"]))
            else:
                j = random.randint(0, total_vectors-1)
                if j < k:
                    reservoir[j] = (vector["id"], vector["dense_values"])
        file_count += 1

    # Prepare query data
    query_ids = [item[0] for item in reservoir]
    query_vectors = np.array([item[1] for item in reservoir], dtype=np.float32)
    
    # Initialize heaps for each query (min-heap for top-k)
    heaps = [[] for _ in range(len(query_vectors))]
    
    # Pass 2: Process dataset in chunks to find top matches
    file_count = 0
    while True:
        path = os.path.join(dataset_dir, f"test{file_count}.parquet")
        if not os.path.exists(path):
            break
            
        df = pd.read_parquet(path)
        chunk_vectors = []
        chunk_ids = []
        
        for index, row in df.iterrows():
            id_val = row[dataset_config["id"]] if dataset_config["id"] else index
            embedding = row[dataset_config["embeddings"]]