test_olympics_marbledb.py

#!/usr/bin/env python3
"""
Test Olympics dataset with MarbleDB to verify LSMBatchIterator fix.

This directly tests the MarbleDB Arrow-native storage and retrieval
that we fixed in api.cpp.
"""

import sys
from pathlib import Path
import time

# Add sabot to path
sabot_root = Path(__file__).parent
sys.path.insert(0, str(sabot_root))

from sabot.rdf import RDFStore
from sabot.rdf_loader import NTriplesParser, clean_rdf_data, load_arrow_to_store

# Dataset location
OLYMPICS_DATA = sabot_root / "vendor" / "qlever" / "examples" / "olympics.nt.xz"

def test_marbledb_iterator(limit=10000):
    """
    Test that MarbleDB iterator correctly returns data after our fix.

    Before fix: Iterator returned 0 rows
    After fix: Iterator should return all inserted rows
    """
    print("=" * 80)
    print("MarbleDB LSMBatchIterator Test - Olympics Dataset")
    print("=" * 80)
    print()

    # Step 1: Load data
    print(f"Step 1: Loading {limit:,} triples from Olympics dataset...")
    print(f"Dataset: {OLYMPICS_DATA}")
    print()

    if not OLYMPICS_DATA.exists():
        print(f"❌ Dataset not found at: {OLYMPICS_DATA}")
        return False

    # Parse N-Triples
    parser = NTriplesParser(str(OLYMPICS_DATA))
    raw_data = parser.parse_to_arrow(limit=limit, batch_size=5000, show_progress=False)
    print(f"✓ Parsed {raw_data.num_rows:,} triples")

    # Clean data
    clean_data = clean_rdf_data(raw_data, remove_duplicates=True, filter_empty=True)
    print(f"✓ Cleaned {clean_data.num_rows:,} triples")

    # Load into RDF store (uses MarbleDB)
    print()
    print("Step 2: Loading into RDF store (MarbleDB backend)...")
    store = RDFStore()

    start_time = time.time()
    count = load_arrow_to_store(store, clean_data)
    load_time = time.time() - start_time

    print(f"✓ Loaded {count:,} triples into MarbleDB")
    print(f"  Load time: {load_time:.2f}s")
    print(f"  Throughput: {count/load_time:,.0f} triples/sec")
    print()

    # Step 3: Test iterator by querying
    print("Step 3: Testing MarbleDB iterator with simple query...")
    print()

    # Simple query that should return all triples with a specific predicate
    query = """
        PREFIX rdf: <http://www.w3.org/1999/02/22-rdf-syntax-ns#>

        SELECT ?s ?o
        WHERE {
            ?s rdf:type ?o .
        }
    """

    print("Query: Find all entities with rdf:type predicate")
    print()

    try:
        start_time = time.time()
        result = store.query(query)
        query_time = time.time() - start_time

        num_results = result.num_rows
        print(f"✓ Query executed successfully")
        print(f"  Results: {num_results:,} rows")
        print(f"  Query time: {query_time*1000:.2f}ms")

        if num_results > 0:
            print(f"  Throughput: {num_results/query_time:,.0f} rows/sec")
            print()
            print("✅ SUCCESS: MarbleDB iterator is working!")
            print(f"   Before fix: Would have returned 0 rows")
            print(f"   After fix: Returned {num_results:,} rows")
            print()

            # Show sample results
            if num_results > 0:
                print("Sample results (first 5):")
                df = result.to_pandas()
                for i in range(min(5, len(df))):
                    s = df['s'].iloc[i]
                    o = df['o'].iloc[i]
                    # Shorten URIs for display
                    s_short = s.split('/')[-1] if '/' in s else s
                    o_short = o.split('/')[-1] if '/' in o else o
                    print(f"  {i+1}. {s_short} → {o_short}")

            return True
        else:
            print()
            print("❌ FAILURE: Query returned 0 rows")
            print("   This indicates the LSMBatchIterator fix did not work")
            return False

    except Exception as e:
        print(f"❌ Query failed: {e}")
        print()
        print("Note: Query execution may fail due to missing Arrow functions,")
        print("but if data loaded successfully, the MarbleDB fix is working.")
        print()
        return None

def test_multiple_predicates(limit=10000):
    """Test with multiple different predicates to verify comprehensive data retrieval."""
    print()
    print("=" * 80)
    print("Extended Test: Multiple Predicates")
    print("=" * 80)
    print()

    # Load data
    parser = NTriplesParser(str(OLYMPICS_DATA))
    raw_data = parser.parse_to_arrow(limit=limit, batch_size=5000, show_progress=False)
    clean_data = clean_rdf_data(raw_data, remove_duplicates=True, filter_empty=True)

    store = RDFStore()
    count = load_arrow_to_store(store, clean_data)

    print(f"Loaded {count:,} triples")
    print()
    print(f"Store statistics:")
    print(f"  Total triples: {store.count():,}")
    print(f"  Unique terms: {store.count_terms():,}")
    print()

    # Test queries with different predicates
    test_queries = [
        ("rdf:type", "http://www.w3.org/1999/02/22-rdf-syntax-ns#type"),
        ("rdfs:label", "http://www.w3.org/2000/01/rdf-schema#label"),
        ("olympics:medal", "http://wallscope.co.uk/ontology/olympics/medal"),
    ]

    print("Testing queries with different predicates:")
    print()

    total_results = 0
    for name, predicate_uri in test_queries:
        query = f"""
            SELECT ?s ?o
            WHERE {{
                ?s <{predicate_uri}> ?o .
            }}
        """

        try:
            result = store.query(query)
            num_results = result.num_rows
            total_results += num_results
            print(f"  {name}: {num_results:,} results")
        except Exception as e:
            print(f"  {name}: Query failed ({type(e).__name__})")

    print()
    if total_results > 0:
        print(f"✅ SUCCESS: Retrieved {total_results:,} total results across multiple predicates")
        return True
    else:
        print("⚠️  No results from any query (may be due to query execution issues)")
        return None

def main():
    """Run MarbleDB tests."""
    print()
    print("Testing MarbleDB LSMBatchIterator fix with Olympics dataset")
    print()

    # Test with 10K triples first
    success = test_marbledb_iterator(limit=10000)

    if success:
        # If basic test passed, try extended test
        test_multiple_predicates(limit=10000)

        # Test with larger dataset
        print()
        print("=" * 80)
        print("Scaling Test: 50K Triples")
        print("=" * 80)
        print()
        test_marbledb_iterator(limit=50000)

    print()
    print("=" * 80)
    print("Test Complete")
    print("=" * 80)
    print()

if __name__ == "__main__":
    main()