cassandra-trunk/src/java/org/apache/cassandra/cql3/functions/CollectionFcts.java
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.cassandra.cql3.functions;

import java.nio.ByteBuffer;
import java.util.List;
import java.util.Map;
import java.util.Set;

import com.google.common.collect.ImmutableList;

import org.apache.cassandra.cql3.CQL3Type;
import org.apache.cassandra.db.marshal.AbstractType;
import org.apache.cassandra.db.marshal.CollectionType;
import org.apache.cassandra.db.marshal.Int32Type;
import org.apache.cassandra.db.marshal.ListType;
import org.apache.cassandra.db.marshal.MapType;
import org.apache.cassandra.db.marshal.SetType;
import org.apache.cassandra.transport.ProtocolVersion;

/**
 * Native CQL functions for collections (sets, list and maps).
 * <p>
 * All the functions provided here are {@link NativeScalarFunction}, and they are meant to be applied to single
 * collection values to perform some kind of aggregation with the elements of the collection argument. When possible,
 * the implementation of these aggregation functions is based on the accross-rows aggregation functions available on
 * {@link AggregateFcts}, so both across-rows and within-collection aggregations have the same behaviour.
 */
public class CollectionFcts
{
    public static void addFunctionsTo(NativeFunctions functions)
    {
        functions.add(new FunctionFactory("map_keys", FunctionParameter.anyMap())
        {
            @Override
            protected NativeFunction doGetOrCreateFunction(List<AbstractType<?>> argTypes, AbstractType<?> receiverType)
            {
                return makeMapKeysFunction(name.name, (MapType<?, ?>) argTypes.get(0));
            }
        });

        functions.add(new FunctionFactory("map_values", FunctionParameter.anyMap())
        {
            @Override
            protected NativeFunction doGetOrCreateFunction(List<AbstractType<?>> argTypes, AbstractType<?> receiverType)
            {
                return makeMapValuesFunction(name.name, (MapType<?, ?>) argTypes.get(0));
            }
        });

        functions.add(new FunctionFactory("collection_count", FunctionParameter.anyCollection())
        {
            @Override
            protected NativeFunction doGetOrCreateFunction(List<AbstractType<?>> argTypes, AbstractType<?> receiverType)
            {
                return makeCollectionCountFunction(name.name, (CollectionType<?>) argTypes.get(0));
            }
        });

        functions.add(new FunctionFactory("collection_min", FunctionParameter.setOrList())
        {
            @Override
            protected NativeFunction doGetOrCreateFunction(List<AbstractType<?>> argTypes, AbstractType<?> receiverType)
            {
                return makeCollectionMinFunction(name.name, (CollectionType<?>) argTypes.get(0));
            }
        });

        functions.add(new FunctionFactory("collection_max", FunctionParameter.setOrList())
        {
            @Override
            protected NativeFunction doGetOrCreateFunction(List<AbstractType<?>> argTypes, AbstractType<?> receiverType)
            {
                return makeCollectionMaxFunction(name.name, (CollectionType<?>) argTypes.get(0));
            }
        });

        functions.add(new FunctionFactory("collection_sum", FunctionParameter.numericSetOrList())
        {
            @Override
            protected NativeFunction doGetOrCreateFunction(List<AbstractType<?>> argTypes, AbstractType<?> receiverType)
            {
                return makeCollectionSumFunction(name.name, (CollectionType<?>) argTypes.get(0));
            }
        });

        functions.add(new FunctionFactory("collection_avg", FunctionParameter.numericSetOrList())
        {
            @Override
            protected NativeFunction doGetOrCreateFunction(List<AbstractType<?>> argTypes, AbstractType<?> receiverType)
            {
                return makeCollectionAvgFunction(name.name, (CollectionType<?>) argTypes.get(0));
            }
        });
    }

    /**
     * Returns a native scalar function for getting the keys of a map column, as a set.
     *
     * @param name      the name of the function
     * @param inputType the type of the map argument by the returned function
     * @param <K>       the class of the map argument keys
     * @param <V>       the class of the map argument values
     * @return a function returning a serialized set containing the keys of the map passed as argument
     */
    private static <K, V> NativeScalarFunction makeMapKeysFunction(String name, MapType<K, V> inputType)
    {
        SetType<K> outputType = SetType.getInstance(inputType.getKeysType(), false);

        return new NativeScalarFunction(name, outputType, inputType)
        {
            @Override
            public ByteBuffer execute(Arguments arguments)
            {
                if (arguments.containsNulls())
                    return null;

                Map<K, V> map = arguments.get(0);
                Set<K> keys = map.keySet();
                return outputType.decompose(keys);
            }
        };
    }

    /**
     * Returns a native scalar function for getting the values of a map column, as a list.
     *
     * @param name      the name of the function
     * @param inputType the type of the map argument accepted by the returned function
     * @param <K>       the class of the map argument keys
     * @param <V>       the class of the map argument values
     * @return a function returning a serialized list containing the values of the map passed as argument
     */
    private static <K, V> NativeScalarFunction makeMapValuesFunction(String name, MapType<K, V> inputType)
    {
        ListType<V> outputType = ListType.getInstance(inputType.getValuesType(), false);

        return new NativeScalarFunction(name, outputType, inputType)
        {
            @Override
            public ByteBuffer execute(Arguments arguments)
            {
                if (arguments.containsNulls())
                    return null;

                Map<K, V> map = arguments.get(0);
                List<V> values = ImmutableList.copyOf(map.values());
                return outputType.decompose(values);
            }
        };
    }

    /**
     * Returns a native scalar function for getting the number of elements in a collection.
     *
     * @param name      the name of the function
     * @param inputType the type of the collection argument accepted by the returned function
     * @param <T>       the type of the elements of the collection argument
     * @return a function returning the number of elements in the collection passed as argument, as a 32-bit integer
     */
    private static <T> NativeScalarFunction makeCollectionCountFunction(String name, CollectionType<T> inputType)
    {
        return new NativeScalarFunction(name, Int32Type.instance, inputType)
        {
            @Override
            public Arguments newArguments(ProtocolVersion version)
            {
                return FunctionArguments.newNoopInstance(version, 1);
            }

            @Override
            public ByteBuffer execute(Arguments arguments)
            {
                if (arguments.containsNulls())
                    return null;


                int size = inputType.size(arguments.get(0));
                return Int32Type.instance.decompose(size);
            }
        };
    }

    /**
     * Returns a native scalar function for getting the min element in a collection.
     *
     * @param name      the name of the function
     * @param inputType the type of the collection argument accepted by the returned function
     * @param <T>       the type of the elements of the collection argument
     * @return a function returning the min element in the collection passed as argument
     */
    private static <T> NativeScalarFunction makeCollectionMinFunction(String name, CollectionType<T> inputType)
    {
        AbstractType<?> elementsType = elementsType(inputType);
        NativeAggregateFunction function = elementsType.isCounter()
                                           ? AggregateFcts.minFunctionForCounter
                                           : AggregateFcts.makeMinFunction(elementsType);
        return new CollectionAggregationFunction(name, inputType, function);
    }

    /**
     * Returns a native scalar function for getting the max element in a collection.
     *
     * @param name      the name of the function
     * @param inputType the type of the collection argument accepted by the returned function
     * @param <T>       the type of the elements of the collection argument
     * @return a function returning the max element in the collection passed as argument
     */
    private static <T> NativeScalarFunction makeCollectionMaxFunction(String name, CollectionType<T> inputType)
    {
        AbstractType<?> elementsType = elementsType(inputType);
        NativeAggregateFunction function = elementsType.isCounter()
                                           ? AggregateFcts.maxFunctionForCounter
                                           : AggregateFcts.makeMaxFunction(elementsType);
        return new CollectionAggregationFunction(name, inputType, function);
    }

    /**
     * Returns a native scalar function for getting the sum of the elements in a numeric collection.
     * </p>
     * The value returned by the function is of the same type as elements of its input collection, so there is a risk
     * of overflow if the sum of the values exceeds the maximum value that the type can represent.
     *
     * @param name      the name of the function
     * @param inputType the type of the collection argument accepted by the returned function
     * @param <T>       the type of the elements of the collection argument
     * @return a function returning the sum of the elements in the collection passed as argument
     */
    private static <T> NativeScalarFunction makeCollectionSumFunction(String name, CollectionType<T> inputType)
    {
        CQL3Type elementsType = elementsType(inputType).asCQL3Type();
        NativeAggregateFunction function = getSumFunction((CQL3Type.Native) elementsType);
        return new CollectionAggregationFunction(name, inputType, function);
    }

    private static NativeAggregateFunction getSumFunction(CQL3Type.Native type)
    {
        switch (type)
        {
            case TINYINT:
                return AggregateFcts.sumFunctionForByte;
            case SMALLINT:
                return AggregateFcts.sumFunctionForShort;
            case INT:
                return AggregateFcts.sumFunctionForInt32;
            case BIGINT:
                return AggregateFcts.sumFunctionForLong;
            case FLOAT:
                return AggregateFcts.sumFunctionForFloat;
            case DOUBLE:
                return AggregateFcts.sumFunctionForDouble;
            case VARINT:
                return AggregateFcts.sumFunctionForVarint;
            case DECIMAL:
                return AggregateFcts.sumFunctionForDecimal;
            default:
                throw new AssertionError("Expected numeric collection but found " + type);
        }
    }

    /**
     * Returns a native scalar function for getting the average of the elements in a numeric collection.
     * </p>
     * The average of an empty collection returns zero. The value returned by the function is of the same type as the
     * elements of its input collection, so if those don't have a decimal part then the returned average won't have a
     * decimal part either.
     *
     * @param name      the name of the function
     * @param inputType the type of the collection argument accepted by the returned function
     * @param <T>       the type of the elements of the collection argument
     * @return a function returning the average value of the elements in the collection passed as argument
     */
    private static <T> NativeScalarFunction makeCollectionAvgFunction(String name, CollectionType<T> inputType)
    {
        CQL3Type elementsType = elementsType(inputType).asCQL3Type();
        NativeAggregateFunction function = getAvgFunction((CQL3Type.Native) elementsType);
        return new CollectionAggregationFunction(name, inputType, function);
    }

    private static NativeAggregateFunction getAvgFunction(CQL3Type.Native type)
    {
        switch (type)
        {
            case TINYINT:
                return AggregateFcts.avgFunctionForByte;
            case SMALLINT:
                return AggregateFcts.avgFunctionForShort;
            case INT:
                return AggregateFcts.avgFunctionForInt32;
            case BIGINT:
                return AggregateFcts.avgFunctionForLong;
            case FLOAT:
                return AggregateFcts.avgFunctionForFloat;
            case DOUBLE:
                return AggregateFcts.avgFunctionForDouble;
            case VARINT:
                return AggregateFcts.avgFunctionForVarint;
            case DECIMAL:
                return AggregateFcts.avgFunctionForDecimal;
            default:
                throw new AssertionError("Expected numeric collection but found " + type);
        }
    }

    /**
     * @return the type of the elements of the specified collection type.
     */
    private static AbstractType<?> elementsType(CollectionType<?> type)
    {
        if (type.kind == CollectionType.Kind.LIST)
        {
            return ((ListType<?>) type).getElementsType();
        }

        if (type.kind == CollectionType.Kind.SET)
        {
            return ((SetType<?>) type).getElementsType();
        }

        throw new AssertionError("Cannot get the element type of: " + type);
    }

    /**
     * A {@link NativeScalarFunction} for aggregating the elements of a collection according to the aggregator of
     * a certain {@link NativeAggregateFunction}.
     * <p>
     * {@link NativeAggregateFunction} is meant to be used for aggregating values accross rows, but here we use that
     * function to aggregate the elements of a single collection value. That way, functions such as {@code avg} and
     * {@code collection_avg} should have the same behaviour when applied to row columns or collection elements.
     */
    private static class CollectionAggregationFunction extends NativeScalarFunction
    {
        private final CollectionType<?> inputType;
        private final NativeAggregateFunction aggregateFunction;

        public CollectionAggregationFunction(String name,
                                             CollectionType<?> inputType,
                                             NativeAggregateFunction aggregateFunction)
        {
            super(name, aggregateFunction.returnType, inputType);
            this.inputType = inputType;
            this.aggregateFunction = aggregateFunction;
        }

        @Override
        public Arguments newArguments(ProtocolVersion version)
        {
            return FunctionArguments.newNoopInstance(version, 1);
        }

        @Override
        public ByteBuffer execute(Arguments arguments)
        {
            if (arguments.containsNulls())
                return null;

            Arguments args = aggregateFunction.newArguments(arguments.getProtocolVersion());
            AggregateFunction.Aggregate aggregate = aggregateFunction.newAggregate();

            inputType.forEach(arguments.get(0), element -> {
                args.set(0, element);
                aggregate.addInput(args);
            });

            return aggregate.compute(arguments.getProtocolVersion());
        }
    }
}
