Semantic Code Search (Java)

Search a Java codebase in natural language — "find the function that validates a JWT token" — instead of grepping for identifiers. AST-based chunking with JavaParser, transformer embeddings via DJL, Qdrant for vector storage, and hybrid dense + BM25 retrieval fused with Reciprocal Rank Fusion.

This is the Java entry in a series of embedding + vector-database projects (the others are in Python); it deliberately shows the same engineering patterns — clean interfaces, a testable core, a measurable design — in a statically typed, build-tooled JVM setting.

What this project demonstrates

• AST-based chunking, not line windows — JavaParser splits each source file along real language boundaries (methods, constructors, classes, interfaces, enums, records). Every embedded vector is a complete construct carrying its signature, Javadoc and exact line range, which is what makes natural-language → code matching actually work.
• A code-aware tokenizer — keyword search over code fails with a naive tokenizer because query words are welded into identifiers. CodeTokenizer splits validateJWTToken → validate, jwt, token (camelCase, snake_case and acronym boundaries), so BM25 can match them.
• Hybrid retrieval — dense embedding search (paraphrase-friendly) fused with BM25 keyword search (exact-identifier-friendly) via Reciprocal Rank Fusion, the same scale-free fusion used in production hybrid search.
• Hexagonal, testable design — Embedder, VectorStore and CorpusRepository are interfaces returning domain types; the ranking logic in SearchService is unit-tested with in-memory fakes, with no model download and no running Qdrant.
• Runs offline for a smoke test — EMBEDDER=hashing swaps the transformer for a deterministic hashing embedder so the whole index → search flow works without network access.

Architecture

  source tree (*.java)
        │  SourceWalker
        ▼
  JavaSourceParser  ──►  List<CodeUnit>   (method / class / … + signature,
        │                                   javadoc, line range, source)
        │  embeddingText() = signature + javadoc + body
        ▼
   Embedder (DJL transformer  |  offline HashingEmbedder)
        │  float[] per unit
        ▼
   Qdrant collection "code"  ◄─────────────┐
   (cosine similarity)                     │  vectors
        ▲                                  │
        │ query vector                     │
        │                          JsonlCorpusRepository
        │                          (canonical CodeUnit records:
        │                           id → unit, and BM25 corpus)
        │                                  │
   ┌────┴───────── SearchService ──────────┘
   │   dense:  Qdrant ANN search
   │   hybrid: Qdrant ANN  ⊕  BM25(CodeTokenizer)  ── Reciprocal Rank Fusion
   ▼
  ranked SearchResults  (printed by the CLI)

Quickstart

Requirements: JDK 17+, Maven, Docker. First run with EMBEDDER=djl downloads the embedding model once (a few hundred MB).

# 1. Build + run the unit tests
mvn verify

# 2. Start Qdrant (the Java client talks gRPC on 6334)
docker compose up -d          # dashboard: http://localhost:6333/dashboard

# 3a. Offline smoke test — no model download, lexical embedder
EMBEDDER=hashing java -jar target/semantic-code-search.jar index sample-src
EMBEDDER=hashing java -jar target/semantic-code-search.jar search "validate a jwt token" --hybrid

# 3b. Real semantic search — transformer embeddings via DJL
java -jar target/semantic-code-search.jar index sample-src
java -jar target/semantic-code-search.jar search "check whether a token is valid"
java -jar target/semantic-code-search.jar search "greatest common divisor" -k 5

Point index at any Java project (… index /path/to/some/repo/src) to search a real codebase.

CLI

index  <source-dir>                    build the index from a source tree
search "<query>" [--hybrid] [-k N]     search (dense by default; --hybrid adds BM25+RRF)

Configuration is via environment variables (see .env.example): QDRANT_HOST, QDRANT_PORT (gRPC, default 6334), COLLECTION, EMBEDDER (djl | hashing), DJL_MODEL_URL, CORPUS_PATH.

Dense vs. hybrid — why both

Dense embeddings excel at paraphrase: "check whether a token is valid" finds validateJwtToken even though no word matches. But embeddings can miss exact lexical hits — a rare method name, an annotation, an error code. BM25 over the code-aware tokenizer catches those. Reciprocal Rank Fusion merges the two rankings using only ranks (score = Σ 1/(k + rank), k=60), which avoids the incompatible-scale problem of adding cosine and BM25 scores directly. The bundled SearchServiceTest demonstrates hybrid rescuing an exact-keyword match that the (fake) dense ranker put second.

Using a code-specific embedding model

The default is all-MiniLM-L6-v2 (general-purpose, reliable, 384-dim). For stronger code retrieval, set DJL_MODEL_URL to a code embedding checkpoint (e.g. a jina-embeddings code model) — nothing else changes, since everything depends on the Embedder interface and the collection is sized from embedder.dimension() at index time.

Project structure

semantic-code-search/
├── pom.xml                         # JavaParser, DJL, Qdrant, Jackson, JUnit; shade fat-jar
├── docker-compose.yml              # local Qdrant (gRPC on 6334)
├── sample-src/                     # tiny sample codebase to index for the demo
├── src/main/java/dev/codesearch/
│   ├── App.java                    # CLI: index | search
│   ├── config/Config.java          # env-based configuration
│   ├── model/                      # CodeUnit, SearchResult, Scored, UnitType (records)
│   ├── ingest/
│   │   ├── SourceWalker.java        # find .java files
│   │   ├── JavaSourceParser.java    # AST -> CodeUnit (the headline feature)
│   │   └── IndexingPipeline.java    # walk -> parse -> embed -> index -> persist
│   ├── embed/
│   │   ├── Embedder.java            # interface
│   │   ├── DjlEmbedder.java         # transformer via DJL
│   │   └── HashingEmbedder.java     # deterministic offline embedder
│   ├── store/
│   │   ├── VectorStore.java         # interface (returns domain types)
│   │   └── QdrantVectorStore.java   # Qdrant gRPC implementation
│   ├── corpus/
│   │   ├── CorpusRepository.java    # interface
│   │   └── JsonlCorpusRepository.java  # JSONL sidecar (Jackson)
│   └── search/
│       ├── CodeTokenizer.java       # camelCase / snake_case / acronym splitting
│       ├── Bm25Index.java           # keyword scoring
│       ├── ReciprocalRankFusion.java
│       └── SearchService.java       # dense + hybrid, pure & unit-testable
├── src/test/java/dev/codesearch/    # JUnit 5: tokenizer, BM25, RRF, service, parser
└── docs/DESIGN_DECISIONS.md

Design decisions

See docs/DESIGN_DECISIONS.md — AST chunking vs. line windows, why the vector store returns domain types, the corpus sidecar alongside Qdrant, string-id → UUID mapping, and the offline embedder.

Limitations & possible extensions

• Java only — the parser is Java-specific; the JavaSourceParser could be generalised behind a SourceParser interface with tree-sitter backends for other languages.
• Batch index, no incremental update — re-indexing rebuilds the collection; a file-hash check to re-embed only changed files is a natural next step.
• No cross-encoder re-ranking — a second-stage re-ranker over the fused candidates would raise precision (mirrors the re-ranker in the companion RAG project).
• In-memory BM25 — fine for typical repositories; a very large corpus would want a persistent sparse index (e.g. Lucene).

License

MIT