<commit_before>package com.kryptnostic.rhizome.mapstores.cassandra;

import java.util.Collection;
import java.util.Map;
import java.util.Map.Entry;
import java.util.stream.Collectors;

import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.tuple.Pair;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.datastax.driver.core.BoundStatement;
import com.datastax.driver.core.PreparedStatement;
import com.datastax.driver.core.RegularStatement;
import com.datastax.driver.core.ResultSet;
import com.datastax.driver.core.ResultSetFuture;
import com.datastax.driver.core.Row;
import com.datastax.driver.core.Session;
import com.datastax.driver.core.exceptions.NoHostAvailableException;
import com.datastax.driver.core.exceptions.QueryExecutionException;
import com.datastax.driver.core.exceptions.QueryValidationException;
import com.google.common.base.Preconditions;
import com.google.common.collect.Iterables;
import com.hazelcast.config.MapConfig;
import com.hazelcast.config.MapStoreConfig;
import com.kryptnostic.rhizome.cassandra.CassandraTableBuilder;
import com.kryptnostic.rhizome.mapstores.TestableSelfRegisteringMapStore;

public abstract class AbstractStructuredCassandraMapstoreBase<K, V> implements TestableSelfRegisteringMapStore<K, V> {
    private static final Logger         logger = LoggerFactory
            .getLogger( AbstractStructuredCassandraMapstoreBase.class );
    private final Session               session;
    private final String                mapName;
    protected final CassandraTableBuilder tableBuilder;
    private final PreparedStatement     allKeysQuery;
    private final PreparedStatement     loadQuery;
    private final PreparedStatement     storeQuery;
    private final PreparedStatement     deleteQuery;

    public AbstractStructuredCassandraMapstoreBase(
            String mapName,
            Session session,
            CassandraTableBuilder tableBuilder ) {
        Preconditions.checkArgument( StringUtils.isNotBlank( mapName ), "Map name cannot be blank" );
        this.mapName = mapName;
        this.session = Preconditions.checkNotNull( session, "Cassandra session cannot be null" );
        this.tableBuilder = Preconditions.checkNotNull( tableBuilder, "Table builder is required" );
        createCassandraSchemaIfNotExist( session, tableBuilder );
        allKeysQuery = prepareLoadAllKeysQuery();
        loadQuery = prepareLoadQuery();
        storeQuery = prepareStoreQuery();
        deleteQuery = prepareDeleteQuery();
    }

    private void createCassandraSchemaIfNotExist(
            Session session,
            CassandraTableBuilder tableBuilder ) {
        session.execute(
                String.format(
                        "CREATE KEYSPACE IF NOT EXISTS %s WITH REPLICATION={ 'class' : 'SimpleStrategy', 'replication_factor' : %d } AND DURABLE_WRITES=true",
                        tableBuilder.getKeyspace().or( "sparks" ),
                        tableBuilder.getReplicationFactor() ) );
        session.execute( tableBuilder.buildCreateTableQuery() );
    }

    @Override
    public V load( K key ) {
        return safeTransform( asyncLoad( key ) );
    }

    @Override
    public Map<K, V> loadAll( Collection<K> keys ) {
        return keys.stream().map( k -> Pair.of( k, asyncLoad( k ) ) )
                .map( p -> Pair.of( p.getLeft(), safeTransform( p.getRight() ) ) )
                .filter( p -> p.getRight() != null )
                .collect( Collectors.toMap( p -> p.getLeft(), p -> p.getRight() ) );
    }

    @Override
    public Iterable<K> loadAllKeys() {
        /*
         * One limitation of this is that if key stream isn't unique then values may get loaded into the map multiple times.
         */
        ResultSet rs = session.execute( getLoadAllKeysQuery().bind() );
        return Iterables.transform( rs, this::mapKey );
    }

    @Override
    public void store( K key, V value ) {<commit_msg>update createCassandraSchemaIfNotExist<commit_after>package com.kryptnostic.rhizome.mapstores.cassandra;

import java.util.Collection;
import java.util.Map;
import java.util.Map.Entry;
import java.util.stream.Collectors;

import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.tuple.Pair;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.datastax.driver.core.BoundStatement;
import com.datastax.driver.core.PreparedStatement;
import com.datastax.driver.core.RegularStatement;
import com.datastax.driver.core.ResultSet;
import com.datastax.driver.core.ResultSetFuture;
import com.datastax.driver.core.Row;
import com.datastax.driver.core.Session;
import com.datastax.driver.core.exceptions.NoHostAvailableException;
import com.datastax.driver.core.exceptions.QueryExecutionException;
import com.datastax.driver.core.exceptions.QueryValidationException;
import com.google.common.base.Preconditions;
import com.google.common.collect.Iterables;
import com.hazelcast.config.MapConfig;
import com.hazelcast.config.MapStoreConfig;
import com.kryptnostic.rhizome.cassandra.CassandraTableBuilder;
import com.kryptnostic.rhizome.mapstores.TestableSelfRegisteringMapStore;

public abstract class AbstractStructuredCassandraMapstoreBase<K, V> implements TestableSelfRegisteringMapStore<K, V> {
    private static final Logger         logger = LoggerFactory
            .getLogger( AbstractStructuredCassandraMapstoreBase.class );
    private final Session               session;
    private final String                mapName;
    protected final CassandraTableBuilder tableBuilder;
    private final PreparedStatement     allKeysQuery;
    private final PreparedStatement     loadQuery;
    private final PreparedStatement     storeQuery;
    private final PreparedStatement     deleteQuery;

    public AbstractStructuredCassandraMapstoreBase(
            String mapName,
            Session session,
            CassandraTableBuilder tableBuilder ) {
        Preconditions.checkArgument( StringUtils.isNotBlank( mapName ), "Map name cannot be blank" );
        this.mapName = mapName;
        this.session = Preconditions.checkNotNull( session, "Cassandra session cannot be null" );
        this.tableBuilder = Preconditions.checkNotNull( tableBuilder, "Table builder is required" );
        createCassandraSchemaIfNotExist( session, tableBuilder );
        allKeysQuery = prepareLoadAllKeysQuery();
        loadQuery = prepareLoadQuery();
        storeQuery = prepareStoreQuery();
        deleteQuery = prepareDeleteQuery();
    }

    private void createCassandraSchemaIfNotExist(
            Session session,
            CassandraTableBuilder tableBuilder ) {
        session.execute(
                String.format(
                        "CREATE KEYSPACE IF NOT EXISTS %s WITH REPLICATION={ 'class' : 'SimpleStrategy', 'replication_factor' : %d } AND DURABLE_WRITES=true",
                        tableBuilder.getKeyspace().or( "sparks" ),
                        tableBuilder.getReplicationFactor() ) );
        tableBuilder.build().forEach( stmt -> session.execute( stmt ) );
    }

    @Override
    public V load( K key ) {
        return safeTransform( asyncLoad( key ) );
    }

    @Override
    public Map<K, V> loadAll( Collection<K> keys ) {
        return keys.stream().map( k -> Pair.of( k, asyncLoad( k ) ) )
                .map( p -> Pair.of( p.getLeft(), safeTransform( p.getRight() ) ) )
                .filter( p -> p.getRight() != null )
                .collect( Collectors.toMap( p -> p.getLeft(), p -> p.getRight() ) );
    }

    @Override
    public Iterable<K> loadAllKeys() {
        /*
         * One limitation of this is that if key stream isn't unique then values may get loaded into the map multiple times.
         */
        ResultSet rs = session.execute( getLoadAllKeysQuery().bind() );
        return Iterables.transform( rs, this::mapKey );
    }

    @Override
    public void store( K key, V value ) {
{'fraction_non_alphanumeric': 0.09171187082124385, 'fraction_numerical': 0.0010092090324208402, 'mean_word_length': 3.1638655462184873, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 24, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '425244', 'n_tokens_mistral': 2217, 'n_tokens_neox': 2163, 'n_words': 501}