Ignite Calcite queries with partitions

CalciteQueryTest.java

package test;

import static org.junit.jupiter.api.Assertions.assertEquals;

import java.util.List;
import java.util.Set;

import org.apache.ignite.Ignite;
import org.apache.ignite.IgniteCache;
import org.apache.ignite.Ignition;
import org.apache.ignite.cache.QueryEntity;
import org.apache.ignite.cache.query.SqlFieldsQuery;
import org.apache.ignite.calcite.CalciteQueryEngineConfiguration;
import org.apache.ignite.configuration.CacheConfiguration;
import org.apache.ignite.configuration.IgniteConfiguration;
import org.apache.ignite.configuration.QueryEngineConfiguration;
import org.apache.ignite.indexing.IndexingQueryEngineConfiguration;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.ValueSource;

public class CalciteQueryTest {

    @ParameterizedTest
    @ValueSource(classes = { IndexingQueryEngineConfiguration.class, CalciteQueryEngineConfiguration.class })
    public void test(Class<? extends QueryEngineConfiguration> queryCfg) throws Exception {

        // Create the qryEngineCfg (h2 / calcite)
        IgniteConfiguration cfg = new IgniteConfiguration();
        cfg.getSqlConfiguration().setQueryEnginesConfiguration(queryCfg.getConstructor().newInstance().setDefault(true));

        try (Ignite ignite = Ignition.start(cfg)) {

            // Create a basic cfg
            CacheConfiguration<String, Foo> cCfg = new CacheConfiguration<>("test");
            QueryEntity qe = new QueryEntity(String.class, Foo.class);
            qe.setTableName("foo");
            qe.addQueryField("id", String.class.getName(), null);
            cCfg.setQueryEntities(Set.of(qe));

            // Insert some values
            IgniteCache<String, Foo> cache = ignite.createCache(cCfg);
            cache.put("1", new Foo("1"));
            cache.put("2", new Foo("2"));
            cache.put("3", new Foo("3"));
            cache.put("4", new Foo("4"));

            int count = 0;

            // Do one query per partitions, and sum results.
            // On H2, queries return result only for the specified partitions, hence most queries return empty results an the total count is 4
            // On Calcite, all queries return all the results, hence the total count is 4096
            for (int part : ignite.affinity("test").primaryPartitions(ignite.cluster().localNode())) {

                // Create query and set partition
                SqlFieldsQuery qry = new SqlFieldsQuery("select _key, _val from foo").setPartitions(part);

                List<List<?>> res = cache.withKeepBinary().query(qry).getAll();
                System.out.printf("Part #%d: %s%n", part, res);
                // count results
                count += res.size();
            }

            assertEquals(4, count, "expected 4 results");
        }
    }

    public static class Foo {

        String id;

        public Foo(String id) {
            this.id = id;
        }

    }
}