/*
 * 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.druid.sql.calcite;

import com.google.common.collect.ImmutableList;
import junitparams.JUnitParamsRunner;
import junitparams.Parameters;
import org.apache.druid.java.util.common.Intervals;
import org.apache.druid.java.util.common.granularity.AllGranularity;
import org.apache.druid.java.util.common.granularity.Granularities;
import org.apache.druid.java.util.common.granularity.Granularity;
import org.apache.druid.query.QueryDataSource;
import org.apache.druid.query.TableDataSource;
import org.apache.druid.query.aggregation.CountAggregatorFactory;
import org.apache.druid.query.aggregation.FilteredAggregatorFactory;
import org.apache.druid.query.aggregation.LongMaxAggregatorFactory;
import org.apache.druid.query.aggregation.LongSumAggregatorFactory;
import org.apache.druid.query.aggregation.any.LongAnyAggregatorFactory;
import org.apache.druid.query.aggregation.cardinality.CardinalityAggregatorFactory;
import org.apache.druid.query.aggregation.hyperloglog.HyperUniqueFinalizingPostAggregator;
import org.apache.druid.query.aggregation.post.ArithmeticPostAggregator;
import org.apache.druid.query.aggregation.post.FieldAccessPostAggregator;
import org.apache.druid.query.dimension.DefaultDimensionSpec;
import org.apache.druid.query.expression.TestExprMacroTable;
import org.apache.druid.query.groupby.GroupByQuery;
import org.apache.druid.segment.column.ValueType;
import org.apache.druid.segment.join.JoinType;
import org.apache.druid.segment.virtual.ExpressionVirtualColumn;
import org.apache.druid.sql.calcite.expression.DruidExpression;
import org.apache.druid.sql.calcite.util.CalciteTests;
import org.junit.Test;
import org.junit.runner.RunWith;

import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;

@RunWith(JUnitParamsRunner.class)
public class CalciteCorrelatedQueryTest extends BaseCalciteQueryTest
{

  @Test
  @Parameters(source = QueryContextForJoinProvider.class)
  public void testCorrelatedSubquery(Map<String, Object> queryContext) throws Exception
  {
    cannotVectorize();
    queryContext = withLeftDirectAccessEnabled(queryContext);

    testQuery(
        "select country, ANY_VALUE(\n"
        + "        select avg(\"users\") from (\n"
        + "            select floor(__time to day), count(distinct user) \"users\" from visits f where f.country = visits.country group by 1\n"
        + "        )\n"
        + "     ) as \"DAU\"\n"
        + "from visits \n"
        + "group by 1",
        queryContext,
        ImmutableList.of(
            GroupByQuery.builder()
                        .setDataSource(
                            join(
                                new TableDataSource(CalciteTests.USERVISITDATASOURCE),
                                new QueryDataSource(
                                    GroupByQuery.builder()
                                                .setDataSource(
                                                    GroupByQuery.builder()
                                                                .setDataSource(CalciteTests.USERVISITDATASOURCE)
                                                                .setQuerySegmentSpec(querySegmentSpec(Intervals.ETERNITY))
                                                                .setVirtualColumns(new ExpressionVirtualColumn(
                                                                    "v0",
                                                                    "timestamp_floor(\"__time\",'P1D',null,'UTC')",
                                                                    ValueType.LONG,
                                                                    TestExprMacroTable.INSTANCE
                                                                ))
                                                                .setDimFilter(not(selector("country", null, null)))
                                                                .setDimensions(
                                                                    new DefaultDimensionSpec(
                                                                        "v0",
                                                                        "d0",
                                                                        ValueType.LONG
                                                                    ),
                                                                    new DefaultDimensionSpec(
                                                                        "country",
                                                                        "d1"
                                                                    )
                                                                )
                                                                .setAggregatorSpecs(new CardinalityAggregatorFactory(
                                                                    "a0:a",
                                                                    null,
                                                                    Collections.singletonList(
                                                                        new DefaultDimensionSpec(
                                                                            "user",
                                                                            "user"
                                                                        )),
                                                                    false,
                                                                    true
                                                                ))
                                                                .setPostAggregatorSpecs(Collections.singletonList(new HyperUniqueFinalizingPostAggregator(
                                                                    "a0",
                                                                    "a0:a"
                                                                )))
                                                                .setContext(
                                                                    withTimestampResultContext(
                                                                        queryContext,
                                                                        "d0",
                                                                        Granularities.DAY
                                                                    )
                                                                )
                                                                .setGranularity(new AllGranularity())
                                                                .build()
                                                )
                                                .setQuerySegmentSpec(querySegmentSpec(Intervals.ETERNITY))
                                                .setDimensions(new DefaultDimensionSpec("d1", "_d0"))
                                                .setAggregatorSpecs(
                                                    new LongSumAggregatorFactory("_a0:sum", "a0"),
                                                    new CountAggregatorFactory("_a0:count")
                                                )
                                                .setPostAggregatorSpecs(Collections.singletonList(new ArithmeticPostAggregator(
                                                    "_a0",
                                                    "quotient",
                                                    Arrays
                                                        .asList(
                                                            new FieldAccessPostAggregator(null, "_a0:sum"),
                                                            new FieldAccessPostAggregator(null, "_a0:count")
                                                        )
                                                )))
                                                .setGranularity(new AllGranularity())
                                                .setContext(queryContext)
                                                .build()
                                ),
                                "j0.",
                                equalsCondition(
                                    DruidExpression.fromColumn("country"),
                                    DruidExpression.fromColumn("j0._d0")
                                ),
                                JoinType.LEFT
                            )
                        )
                        .setQuerySegmentSpec(querySegmentSpec(Intervals.ETERNITY))
                        .setDimensions(new DefaultDimensionSpec("country", "d0"))
                        .setAggregatorSpecs(new LongAnyAggregatorFactory("a0", "j0._a0"))
                        .setGranularity(new AllGranularity())
                        .setContext(queryContext)
                        .build()
        ),
        ImmutableList.of(
            new Object[]{"India", 2L},
            new Object[]{"USA", 1L},
            new Object[]{"canada", 3L}
        )
    );
  }

  @Test
  @Parameters(source = QueryContextForJoinProvider.class)
  public void testCorrelatedSubqueryWithLeftFilter(Map<String, Object> queryContext) throws Exception
  {
    cannotVectorize();
    queryContext = withLeftDirectAccessEnabled(queryContext);

    testQuery(
        "select country, ANY_VALUE(\n"
        + "        select max(\"users\") from (\n"
        + "            select floor(__time to day), count(*) \"users\" from visits f where f.country = visits.country group by 1\n"
        + "        )\n"
        + "     ) as \"dailyVisits\"\n"
        + "from visits \n"
        + " where city = 'B' and __time between '2021-01-01 01:00:00' AND '2021-01-02 23:59:59'"
        + " group by 1",
        queryContext,
        ImmutableList.of(
            GroupByQuery.builder()
                        .setDataSource(
                            join(
                                new TableDataSource(CalciteTests.USERVISITDATASOURCE),
                                new QueryDataSource(
                                    GroupByQuery.builder()
                                                .setDataSource(
                                                    GroupByQuery.builder()
                                                                .setDataSource(CalciteTests.USERVISITDATASOURCE)
                                                                .setQuerySegmentSpec(querySegmentSpec(Intervals.ETERNITY))
                                                                .setVirtualColumns(new ExpressionVirtualColumn(
                                                                    "v0",
                                                                    "timestamp_floor(\"__time\",'P1D',null,'UTC')",
                                                                    ValueType.LONG,
                                                                    TestExprMacroTable.INSTANCE
                                                                ))
                                                                .setDimFilter(not(selector("country", null, null)))
                                                                .setDimensions(
                                                                    new DefaultDimensionSpec(
                                                                        "v0",
                                                                        "d0",
                                                                        ValueType.LONG
                                                                    ),
                                                                    new DefaultDimensionSpec(
                                                                        "country",
                                                                        "d1"
                                                                    )
                                                                )
                                                                .setAggregatorSpecs(new CountAggregatorFactory("a0"))
                                                                .setContext(
                                                                    withTimestampResultContext(
                                                                        queryContext,
                                                                        "d0",
                                                                        Granularities.DAY
                                                                    )
                                                                )
                                                                .setGranularity(new AllGranularity())
                                                                .build()
                                                )
                                                .setQuerySegmentSpec(querySegmentSpec(Intervals.ETERNITY))
                                                .setDimensions(new DefaultDimensionSpec("d1", "_d0"))
                                                .setAggregatorSpecs(
                                                    new LongMaxAggregatorFactory("_a0", "a0")
                                                )
                                                .setGranularity(new AllGranularity())
                                                .setContext(queryContext)
                                                .build()
                                ),
                                "j0.",
                                equalsCondition(
                                    DruidExpression.fromColumn("country"),
                                    DruidExpression.fromColumn("j0._d0")
                                ),
                                JoinType.LEFT,
                                selector("city", "B", null)
                            )
                        )
                        .setQuerySegmentSpec(querySegmentSpec(Intervals.of(
                            "2021-01-01T01:00:00.000Z/2021-01-02T23:59:59.001Z")))
                        .setDimensions(new DefaultDimensionSpec("country", "d0"))
                        .setAggregatorSpecs(new LongAnyAggregatorFactory("a0", "j0._a0"))
                        .setGranularity(new AllGranularity())
                        .setContext(queryContext)
                        .build()
        ),
        ImmutableList.of(
            new Object[]{"canada", 4L}
        )
    );
  }

  @Test
  @Parameters(source = QueryContextForJoinProvider.class)
  public void testCorrelatedSubqueryWithLeftFilter_leftDirectAccessDisabled(Map<String, Object> queryContext) throws Exception
  {
    cannotVectorize();

    testQuery(
        "select country, ANY_VALUE(\n"
        + "        select max(\"users\") from (\n"
        + "            select floor(__time to day), count(*) \"users\" from visits f where f.country = visits.country group by 1\n"
        + "        )\n"
        + "     ) as \"dailyVisits\"\n"
        + "from visits \n"
        + " where city = 'B' and __time between '2021-01-01 01:00:00' AND '2021-01-02 23:59:59'"
        + " group by 1",
        queryContext,
        ImmutableList.of(
            GroupByQuery.builder()
                        .setDataSource(
                            join(
                                new QueryDataSource(newScanQueryBuilder().dataSource(CalciteTests.USERVISITDATASOURCE)
                                                                         .intervals(querySegmentSpec(Intervals.of(
                                                                             "2021-01-01T01:00:00.000Z/2021-01-02T23:59:59.001Z")))
                                                                         .filters(selector("city", "B", null))
                                                                         .columns("__time", "city", "country")
                                                                         .build()),
                                new QueryDataSource(
                                    GroupByQuery.builder()
                                                .setDataSource(
                                                    GroupByQuery.builder()
                                                                .setDataSource(CalciteTests.USERVISITDATASOURCE)
                                                                .setQuerySegmentSpec(querySegmentSpec(Intervals.ETERNITY))
                                                                .setVirtualColumns(new ExpressionVirtualColumn(
                                                                    "v0",
                                                                    "timestamp_floor(\"__time\",'P1D',null,'UTC')",
                                                                    ValueType.LONG,
                                                                    TestExprMacroTable.INSTANCE
                                                                ))
                                                                .setDimFilter(not(selector("country", null, null)))
                                                                .setDimensions(
                                                                    new DefaultDimensionSpec(
                                                                        "v0",
                                                                        "d0",
                                                                        ValueType.LONG
                                                                    ),
                                                                    new DefaultDimensionSpec(
                                                                        "country",
                                                                        "d1"
                                                                    )
                                                                )
                                                                .setAggregatorSpecs(new CountAggregatorFactory("a0"))
                                                                .setContext(
                                                                    withTimestampResultContext(
                                                                        queryContext,
                                                                        "d0",
                                                                        Granularities.DAY
                                                                    )
                                                                )
                                                                .setGranularity(new AllGranularity())
                                                                .build()
                                                )
                                                .setQuerySegmentSpec(querySegmentSpec(Intervals.ETERNITY))
                                                .setDimensions(new DefaultDimensionSpec("d1", "_d0"))
                                                .setAggregatorSpecs(
                                                    new LongMaxAggregatorFactory("_a0", "a0")
                                                )
                                                .setGranularity(new AllGranularity())
                                                .setContext(queryContext)
                                                .build()
                                ),
                                "j0.",
                                equalsCondition(
                                    DruidExpression.fromColumn("country"),
                                    DruidExpression.fromColumn("j0._d0")
                                ),
                                JoinType.LEFT
                            )
                        )
                        .setQuerySegmentSpec(querySegmentSpec(Intervals.ETERNITY))
                        .setDimensions(new DefaultDimensionSpec("country", "d0"))
                        .setAggregatorSpecs(new LongAnyAggregatorFactory("a0", "j0._a0"))
                        .setGranularity(new AllGranularity())
                        .setContext(queryContext)
                        .build()
        ),
        ImmutableList.of(
            new Object[]{"canada", 4L}
        )
    );
  }

  @Test
  @Parameters(source = QueryContextForJoinProvider.class)
  public void testCorrelatedSubqueryWithCorrelatedQueryFilter(Map<String, Object> queryContext) throws Exception
  {
    cannotVectorize();
    queryContext = withLeftDirectAccessEnabled(queryContext);

    testQuery(
        "select country, ANY_VALUE(\n"
        + "        select max(\"users\") from (\n"
        + "            select floor(__time to day), count(user) \"users\" from visits f where f.country = visits.country and f.city = 'A' group by 1\n"
        + "        )\n"
        + "     ) as \"dailyVisits\"\n"
        + "from visits \n"
        + " where city = 'B'"
        + " group by 1",
        queryContext,
        ImmutableList.of(
            GroupByQuery.builder()
                        .setDataSource(
                            join(
                                new TableDataSource(CalciteTests.USERVISITDATASOURCE),
                                new QueryDataSource(
                                    GroupByQuery.builder()
                                                .setDataSource(
                                                    GroupByQuery.builder()
                                                                .setDataSource(CalciteTests.USERVISITDATASOURCE)
                                                                .setQuerySegmentSpec(querySegmentSpec(Intervals.ETERNITY))
                                                                .setVirtualColumns(new ExpressionVirtualColumn(
                                                                    "v0",
                                                                    "timestamp_floor(\"__time\",'P1D',null,'UTC')",
                                                                    ValueType.LONG,
                                                                    TestExprMacroTable.INSTANCE
                                                                ))
                                                                .setDimensions(
                                                                    new DefaultDimensionSpec(
                                                                        "v0",
                                                                        "d0",
                                                                        ValueType.LONG
                                                                    ),
                                                                    new DefaultDimensionSpec(
                                                                        "country",
                                                                        "d1"
                                                                    )
                                                                )
                                                                .setAggregatorSpecs(new FilteredAggregatorFactory(
                                                                    new CountAggregatorFactory("a0"),
                                                                    not(selector("user", null, null))
                                                                ))
                                                                .setDimFilter(and(
                                                                    selector("city", "A", null),
                                                                    not(selector("country", null, null))
                                                                ))
                                                                .setContext(
                                                                    withTimestampResultContext(
                                                                        queryContext,
                                                                        "d0",
                                                                        Granularities.DAY
                                                                    )
                                                                )
                                                                .setGranularity(new AllGranularity())
                                                                .build()
                                                )
                                                .setQuerySegmentSpec(querySegmentSpec(Intervals.ETERNITY))
                                                .setDimensions(new DefaultDimensionSpec("d1", "_d0"))
                                                .setAggregatorSpecs(
                                                    new LongMaxAggregatorFactory("_a0", "a0")
                                                )
                                                .setGranularity(new AllGranularity())
                                                .setContext(queryContext)
                                                .build()
                                ),
                                "j0.",
                                equalsCondition(
                                    DruidExpression.fromColumn("country"),
                                    DruidExpression.fromColumn("j0._d0")
                                ),
                                JoinType.LEFT,
                                selector("city", "B", null)
                            )
                        )
                        .setQuerySegmentSpec(querySegmentSpec(Intervals.ETERNITY))
                        .setDimensions(new DefaultDimensionSpec("country", "d0"))
                        .setAggregatorSpecs(new LongAnyAggregatorFactory("a0", "j0._a0"))
                        .setGranularity(new AllGranularity())
                        .setContext(queryContext)
                        .build()
        ),
        ImmutableList.of(
            new Object[]{"canada", 2L}
        )
    );
  }

  @Test
  @Parameters(source = QueryContextForJoinProvider.class)
  public void testCorrelatedSubqueryWithCorrelatedQueryFilter_Scan(Map<String, Object> queryContext) throws Exception
  {
    cannotVectorize();
    queryContext = withLeftDirectAccessEnabled(queryContext);
    testQuery(
        "select country, ANY_VALUE(\n"
        + "        select max(\"users\") from (\n"
        + "            select floor(__time to day), count(user) \"users\" from visits f where f.country = visits.country and f.city = 'A' group by 1\n"
        + "        )\n"
        + "     ) as \"dailyVisits\"\n"
        + "from visits \n"
        + " where city = 'B'"
        + " group by 1",
        queryContext,
        ImmutableList.of(
            GroupByQuery.builder()
                        .setDataSource(
                            join(
                                new TableDataSource(CalciteTests.USERVISITDATASOURCE),
                                new QueryDataSource(
                                    GroupByQuery.builder()
                                                .setDataSource(
                                                    GroupByQuery.builder()
                                                                .setDataSource(CalciteTests.USERVISITDATASOURCE)
                                                                .setQuerySegmentSpec(querySegmentSpec(Intervals.ETERNITY))
                                                                .setVirtualColumns(new ExpressionVirtualColumn(
                                                                    "v0",
                                                                    "timestamp_floor(\"__time\",'P1D',null,'UTC')",
                                                                    ValueType.LONG,
                                                                    TestExprMacroTable.INSTANCE
                                                                ))
                                                                .setDimensions(
                                                                    new DefaultDimensionSpec(
                                                                        "v0",
                                                                        "d0",
                                                                        ValueType.LONG
                                                                    ),
                                                                    new DefaultDimensionSpec(
                                                                        "country",
                                                                        "d1"
                                                                    )
                                                                )
                                                                .setAggregatorSpecs(new FilteredAggregatorFactory(
                                                                    new CountAggregatorFactory("a0"),
                                                                    not(selector("user", null, null))
                                                                ))
                                                                .setDimFilter(and(
                                                                    selector("city", "A", null),
                                                                    not(selector("country", null, null))
                                                                ))
                                                                .setContext(
                                                                    withTimestampResultContext(
                                                                        queryContext,
                                                                        "d0",
                                                                        Granularities.DAY
                                                                    )
                                                                )
                                                                .setGranularity(new AllGranularity())
                                                                .build()
                                                )
                                                .setQuerySegmentSpec(querySegmentSpec(Intervals.ETERNITY))
                                                .setDimensions(new DefaultDimensionSpec("d1", "_d0"))
                                                .setAggregatorSpecs(
                                                    new LongMaxAggregatorFactory("_a0", "a0")
                                                )
                                                .setGranularity(new AllGranularity())
                                                .setContext(queryContext)
                                                .build()
                                ),
                                "j0.",
                                equalsCondition(
                                    DruidExpression.fromColumn("country"),
                                    DruidExpression.fromColumn("j0._d0")
                                ),
                                JoinType.LEFT,
                                selector("city", "B", null)
                            )
                        )
                        .setQuerySegmentSpec(querySegmentSpec(Intervals.ETERNITY))
                        .setDimensions(new DefaultDimensionSpec("country", "d0"))
                        .setAggregatorSpecs(new LongAnyAggregatorFactory("a0", "j0._a0"))
                        .setGranularity(new AllGranularity())
                        .setContext(queryContext)
                        .build()
        ),
        ImmutableList.of(
            new Object[]{"canada", 2L}
        )
    );
  }

  private Map<String, Object> withTimestampResultContext(
      Map<String, Object> input,
      String timestampResultField,
      Granularity granularity
  )
  {
    Map<String, Object> output = new HashMap<>(input);
    output.put(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD, timestampResultField);
    output.put(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD_GRANULARITY, granularity);
    output.put(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD_INDEX, 0);
    return output;
  }
}