<commit_before>package org.semagrow.geospatial.selector;

import org.eclipse.rdf4j.model.IRI;
import org.eclipse.rdf4j.model.Literal;
import org.eclipse.rdf4j.model.Resource;
import org.eclipse.rdf4j.model.ValueFactory;
import org.eclipse.rdf4j.model.impl.SimpleValueFactory;
import org.eclipse.rdf4j.model.vocabulary.GEO;
import org.eclipse.rdf4j.query.BindingSet;
import org.eclipse.rdf4j.query.Dataset;
import org.eclipse.rdf4j.query.algebra.StatementPattern;
import org.eclipse.rdf4j.query.algebra.TupleExpr;
import org.eclipse.rdf4j.query.algebra.ValueExpr;
import org.eclipse.rdf4j.query.algebra.helpers.StatementPatternCollector;
import org.eclipse.rdf4j.query.algebra.helpers.VarNameCollector;
import org.eclipse.rdf4j.query.impl.EmptyBindingSet;
import org.eclipse.rdf4j.repository.Repository;
import org.semagrow.plan.util.BPGCollector;
import org.semagrow.plan.util.FilterCollector;
import org.semagrow.selector.QueryAwareSourceSelector;
import org.semagrow.selector.SourceMetadata;
import org.semagrow.selector.SourceSelector;
import org.semagrow.selector.SourceSelectorWrapper;

import java.util.*;

public class GeospatialSourceSelector extends SourceSelectorWrapper implements QueryAwareSourceSelector  {

    private ValueFactory vf = SimpleValueFactory.getInstance();
    private IRI HAS_GEOMETRY = vf.createIRI(GEO.NAMESPACE, "hasGeometry");

    private BoundingBoxBase boundingBoxBase = new BoundingBoxBase();
    private Map<StatementPattern,Collection<SourceMetadata>> selectorMap = new HashMap<>();
    private boolean processed = false;

    public GeospatialSourceSelector(SourceSelector selector) {
        super(selector);
    }

    public void setMetadataRepository(Repository metadata) {
        boundingBoxBase.setMetadata(metadata);
    }

    @Override
    public void processTupleExpr(TupleExpr expr) {
        if (getWrappedSelector() instanceof QueryAwareSourceSelector) {
            ((QueryAwareSourceSelector) getWrappedSelector()).processTupleExpr(expr);
        }
        Collection<TupleExpr> bpgs = BPGCollector.process(expr);
        for (TupleExpr bpg: bpgs) {
            processBPG(bpg);
        }
        processed = true;
    }

    private void processBPG(TupleExpr expr) {

        Collection<StatementPattern> patterns = StatementPatternCollector.process(expr);
        Collection<ValueExpr> filters = FilterCollector.process(expr);

        for (StatementPattern patternAsWKT: patterns) {
            if (patternAsWKT.getPredicateVar().hasValue() && patternAsWKT.getPredicateVar().getValue().equals(GEO.AS_WKT)) {
                /* we found a ?s geo:asWKT ?o */

                Collection<SourceMetadata> candidateSources = getWrappedSelector().getSources(patternAsWKT, null, EmptyBindingSet.getInstance());
                Set<SourceMetadata> prunedSources = new HashSet<>();

                /* search for corresponding filter */
                for (ValueExpr filter: filters) {
                    Set<String> filterVars = VarNameCollector.process(filter);
                    String varName = patternAsWKT.getObjectVar().getName();

                    if (filterVars.contains(varName) && filterVars.size() == 1) {
                        for (SourceMetadata source: candidateSources) {
                            Literal boundingBox = boundingBoxBase.getDatasetBoundingBox(endpointOfSource(source));
                            if (boundingBox != null) {
                               if (BoundingBoxPruner.prune(filter, varName, boundingBox)) {
                                   prunedSources.add(source);
                               }
                            }
                        }
                    }
                }
                Set<SourceMetadata> sources = new HashSet<>();
                sources.addAll(candidateSources);
                sources.removeAll(prunedSources);

                selectorMap.put(patternAsWKT, sources);

                /* search for corresponding hasGeometry pattern */
                for (StatementPattern patternHasGeometry: patterns) {
                    if (patternHasGeometry.getPredicateVar().hasValue() &&
                            patternHasGeometry.getPredicateVar().getValue().equals(HAS_GEOMETRY) &&
                            patternHasGeometry.getObjectVar().getName().equals(patternAsWKT.getSubjectVar().getName())) {
                        selectorMap.put(patternHasGeometry, sources);
                    }
                }
            }
        }<commit_msg>fix geospatial source selector for hasGeometry patterns<commit_after>package org.semagrow.geospatial.selector;

import org.eclipse.rdf4j.model.IRI;
import org.eclipse.rdf4j.model.Literal;
import org.eclipse.rdf4j.model.Resource;
import org.eclipse.rdf4j.model.ValueFactory;
import org.eclipse.rdf4j.model.impl.SimpleValueFactory;
import org.eclipse.rdf4j.model.vocabulary.GEO;
import org.eclipse.rdf4j.query.BindingSet;
import org.eclipse.rdf4j.query.Dataset;
import org.eclipse.rdf4j.query.algebra.StatementPattern;
import org.eclipse.rdf4j.query.algebra.TupleExpr;
import org.eclipse.rdf4j.query.algebra.ValueExpr;
import org.eclipse.rdf4j.query.algebra.helpers.StatementPatternCollector;
import org.eclipse.rdf4j.query.algebra.helpers.VarNameCollector;
import org.eclipse.rdf4j.query.impl.EmptyBindingSet;
import org.eclipse.rdf4j.repository.Repository;
import org.semagrow.plan.util.BPGCollector;
import org.semagrow.plan.util.FilterCollector;
import org.semagrow.selector.QueryAwareSourceSelector;
import org.semagrow.selector.SourceMetadata;
import org.semagrow.selector.SourceSelector;
import org.semagrow.selector.SourceSelectorWrapper;

import java.util.*;
import java.util.stream.Collectors;

public class GeospatialSourceSelector extends SourceSelectorWrapper implements QueryAwareSourceSelector  {

    private ValueFactory vf = SimpleValueFactory.getInstance();
    private IRI HAS_GEOMETRY = vf.createIRI(GEO.NAMESPACE, "hasGeometry");

    private BoundingBoxBase boundingBoxBase = new BoundingBoxBase();
    private Map<StatementPattern,Collection<SourceMetadata>> selectorMap = new HashMap<>();
    private boolean processed = false;

    public GeospatialSourceSelector(SourceSelector selector) {
        super(selector);
    }

    public void setMetadataRepository(Repository metadata) {
        boundingBoxBase.setMetadata(metadata);
    }

    @Override
    public void processTupleExpr(TupleExpr expr) {
        if (getWrappedSelector() instanceof QueryAwareSourceSelector) {
            ((QueryAwareSourceSelector) getWrappedSelector()).processTupleExpr(expr);
        }
        Collection<TupleExpr> bpgs = BPGCollector.process(expr);
        for (TupleExpr bpg: bpgs) {
            processBPG(bpg);
        }
        processed = true;
    }

    private void processBPG(TupleExpr expr) {

        Collection<StatementPattern> patterns = StatementPatternCollector.process(expr);
        Collection<ValueExpr> filters = FilterCollector.process(expr);
        Map<String,Set<Resource>> hasGeoMap = new HashMap<>();

        for (StatementPattern pattern: patterns) {
            if (pattern.getPredicateVar().hasValue() && pattern.getPredicateVar().getValue().equals(GEO.AS_WKT)) {
                /* we found a ?s geo:asWKT ?o */

                Collection<SourceMetadata> candidateSources = getWrappedSelector().getSources(pattern, null, EmptyBindingSet.getInstance());
                Set<SourceMetadata> prunedSources = new HashSet<>();
                Set<Resource> endpoints = candidateSources.stream().map(s -> endpointOfSource(s)).collect(Collectors.toSet());

                /* search for corresponding filter */
                for (ValueExpr filter: filters) {
                    Set<String> filterVars = VarNameCollector.process(filter);
                    String varName = pattern.getObjectVar().getName();

                    if (filterVars.contains(varName) && filterVars.size() == 1) {
                        for (SourceMetadata source: candidateSources) {
                            Literal boundingBox = boundingBoxBase.getDatasetBoundingBox(endpointOfSource(source));
                            if (boundingBox != null) {
                               if (BoundingBoxPruner.prune(filter, varName, boundingBox)) {
                                   prunedSources.add(source);
                                   endpoints.remove(endpointOfSource(source));
                               }
                            }
                        }
                    }
                }
                Set<SourceMetadata> sources = new HashSet<>();
                sources.addAll(candidateSources);
                sources.removeAll(prunedSources);

                selectorMap.put(pattern, sources);
                hasGeoMap.put(pattern.getSubjectVar().getName(), endpoints);
            }
        }

        /* search for corresponding hasGeometry patterns */
        for (StatementPattern pattern: patterns) {
            if (pattern.getPredicateVar().hasValue() && pattern.getPredicateVar().getValue().equals(HAS_GEOMETRY)) {
                /* we found a ?s geo:hasGeometry ?g */

                Collection<SourceMetadata> candidateSources = getWrappedSelector().getSources(pattern, null, EmptyBindingSet.getInstance());
                Set<SourceMetadata> prunedSources = new HashSet<>();

                for (SourceMetadata source: candidateSources) {
                    if (hasGeoMap.get(pattern.getObjectVar().getName()).contains(endpointOfSource(source))) {
                        prunedSources.add(source);
                    }
                }
                selectorMap.put(pattern, prunedSources);
            }
        }
{'fraction_non_alphanumeric': 0.09648127128263337, 'fraction_numerical': 0.0033020328139510887, 'mean_word_length': 2.6408715251690458, 'pattern_counts': {'":': 0, '<': 34, '<?xml version=': 0, '>': 35, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 5, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3232766', 'n_tokens_mistral': 2620, 'n_tokens_neox': 2487, 'n_words': 428}