Simulate a GIN Index for JSONB in Python

import json
import random
from collections import defaultdict

# Mock GIN (Generalized Inverted Index) for JSONB key-value pairs
class GINIndex:
    def __init__(self):
        self.posting_lists = defaultdict(list)  # token -> list of doc_ids
    
    def index(self, doc_id, json_obj):
        """Index a JSON document by flattening all leaf key-value pairs"""
        def flatten(obj, prefix=""):
            if isinstance(obj, dict):
                for k, v in obj.items():
                    key = f"{prefix}.{k}" if prefix else k
                    if isinstance(v, (dict, list)):
                        flatten(v, key)
                    else:
                        token = f"{key}={v}"
                        self.posting_lists[token].append(doc_id)
            elif isinstance(obj, list):
                for i, v in enumerate(obj):
                    flatten(v, f"{prefix}[{i}]")
        
        flatten(json_obj)
    
    def search(self, key, value):
        """Search documents by key-value pair"""
        token = f"{key}={value}"
        return self.posting_lists.get(token, [])

# Demo usage
if __name__ == "__main__":
    gin = GINIndex()
    
    docs = [
        {"id": 1, "name": "Alice", "age": 30, "address": {"city": "NYC", "zip": 10001}},
        {"id": 2, "name": "Bob", "age": 25, "address": {"city": "LA", "zip": 90001}},
        {"id": 3, "name": "Charlie", "age": 30, "address": {"city": "NYC", "zip": 10002}},
        {"id": 4, "name": "David", "age": 35, "address": {"city": "SF", "zip": 94101}}
    ]
    
    for doc in docs:
        gin.index(doc["id"], doc)
    
    # Query examples
    print("People aged 30:", gin.search("age", 30))
    print("People in NYC:", gin.search("address.city", "NYC"))
    print("People with zip 90001:", gin.search("address.zip", 90001))