Skip to content

Commit

Permalink
Add support for a new relation asOfJoin operation (#3162)
Browse files Browse the repository at this point in the history
  • Loading branch information
pierredebelen authored Oct 8, 2024
1 parent f762886 commit fbcbc42
Show file tree
Hide file tree
Showing 28 changed files with 2,168 additions and 131 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -103,6 +103,7 @@ private Completer(String buildCodeContext, MutableList<CompleterExtension> exten
new PivotHandler(),
new SortHandler(),
new JoinHandler(),
new AsOfJoinHandler(),
new SelectHandler(),
new DistinctHandler(),
new OverHandler()
Expand Down Expand Up @@ -268,7 +269,7 @@ private MutableList<String> getFunctionCandidates(org.finos.legend.pure.m3.corei
if (org.finos.legend.pure.m3.navigation.type.Type.subTypeOf(leftType._rawType(), pureModel.getType(M3Paths.Relation), pureModel.getExecutionSupport().getProcessorSupport()))
{
// May want to assert the mul to 1
return Lists.mutable.with("cast", "distinct", "drop", "select", "extend", "filter", "from", "groupBy", "pivot", "join", "limit", "rename", "size", "slice", "sort");
return Lists.mutable.with("cast", "distinct", "drop", "select", "extend", "filter", "from", "groupBy", "pivot", "join", "asOfJoin", "limit", "rename", "size", "slice", "sort");
}
else if (leftType._rawType().getName().equals("String"))
{
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,83 @@
// Copyright 2024 Goldman Sachs
//
// Licensed 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.finos.legend.engine.repl.autocomplete.handlers;

import org.eclipse.collections.api.list.MutableList;
import org.eclipse.collections.impl.factory.Lists;
import org.finos.legend.engine.language.pure.compiler.toPureGraph.CompileContext;
import org.finos.legend.engine.language.pure.compiler.toPureGraph.ProcessingContext;
import org.finos.legend.engine.language.pure.compiler.toPureGraph.PureModel;
import org.finos.legend.engine.language.pure.compiler.toPureGraph.ValueSpecificationBuilder;
import org.finos.legend.engine.protocol.pure.v1.model.valueSpecification.ValueSpecification;
import org.finos.legend.engine.protocol.pure.v1.model.valueSpecification.Variable;
import org.finos.legend.engine.protocol.pure.v1.model.valueSpecification.application.AppliedFunction;
import org.finos.legend.engine.protocol.pure.v1.model.valueSpecification.raw.Lambda;
import org.finos.legend.engine.repl.autocomplete.Completer;
import org.finos.legend.engine.repl.autocomplete.CompletionItem;
import org.finos.legend.engine.repl.autocomplete.FunctionHandler;
import org.finos.legend.pure.m3.coreinstance.meta.pure.metamodel.type.generics.GenericType;

public class AsOfJoinHandler extends FunctionHandler
{
@Override
public String functionName()
{
return "asOfJoin";
}

@Override
public MutableList<CompletionItem> proposedParameters(AppliedFunction currentFunc, GenericType leftType, PureModel pureModel, Completer completer, ProcessingContext processingContext, ValueSpecification currentVS)
{
if (currentFunc.parameters.size() == 2)
{
return completer.processValueSpecification(currentFunc.parameters.get(1), currentVS, pureModel, processingContext).getCompletion();
}
return Lists.mutable.empty();
}

public void handleFunctionAppliedParameters(AppliedFunction currentFunc, GenericType leftType, ProcessingContext processingContext, PureModel pureModel)
{
if (currentFunc.parameters.size() == 3)
{
processLambda(currentFunc, currentFunc.parameters.get(2), leftType, processingContext, pureModel);
}
if (currentFunc.parameters.size() == 4)
{
processLambda(currentFunc, currentFunc.parameters.get(3), leftType, processingContext, pureModel);
}
}

private static void processLambda(AppliedFunction currentFunc, ValueSpecification value, GenericType leftType, ProcessingContext processingContext, PureModel pureModel)
{
GenericType genericType = currentFunc.parameters.get(1).accept(new ValueSpecificationBuilder(new CompileContext.Builder(pureModel).build(), Lists.mutable.empty(), processingContext))._genericType()._typeArguments().getFirst();
if (value instanceof Lambda)
{
processLambda((Lambda)value, leftType._typeArguments().getFirst(), genericType, processingContext, pureModel);
}
}


private static void processLambda(Lambda lambda, GenericType first, GenericType second, ProcessingContext processingContext, PureModel pureModel)
{
if (lambda != null)
{
Variable f_variable = lambda.parameters.get(0);
processingContext.addInferredVariables(f_variable.name, buildTypedVariable(f_variable, first, pureModel.getMultiplicity("one"), pureModel));
Variable s_variable = lambda.parameters.get(1);
processingContext.addInferredVariables(s_variable.name, buildTypedVariable(s_variable, second, pureModel.getMultiplicity("one"), pureModel));
}
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -47,14 +47,14 @@ public void testAutocompleteFunctionParameter()
@Test
public void testArrowOnFunction()
{
Assert.assertEquals("[cast , cast(], [distinct , distinct(], [drop , drop(], [select , select(], [extend , extend(], [filter , filter(], [from , from(], [groupBy , groupBy(], [pivot , pivot(], [join , join(], [limit , limit(], [rename , rename(], [size , size(], [slice , slice(], [sort , sort(]", checkResultNoException(new Completer("###Relational\nDatabase a::A(Table t(col VARCHAR(200)))", Lists.mutable.with(new RelationalCompleterExtension())).complete("#>{a::A.t}#->filter(x|$x.col == 'oo')->")));
Assert.assertEquals("[cast , cast(], [distinct , distinct(], [drop , drop(], [select , select(], [extend , extend(], [filter , filter(], [from , from(], [groupBy , groupBy(], [pivot , pivot(], [join , join(], [asOfJoin , asOfJoin(], [limit , limit(], [rename , rename(], [size , size(], [slice , slice(], [sort , sort(]", checkResultNoException(new Completer("###Relational\nDatabase a::A(Table t(col VARCHAR(200)))", Lists.mutable.with(new RelationalCompleterExtension())).complete("#>{a::A.t}#->filter(x|$x.col == 'oo')->")));
Assert.assertEquals("PARSER error at [6:1-23]: parsing error", new Completer("###Relational\nDatabase a::A(Table t(col VARCHAR(200)))", Lists.mutable.with(new RelationalCompleterExtension())).complete("#>{a::A.t}#->limit(10)-").getEngineException().toPretty());
}

@Test
public void testArrowRelation()
{
Assert.assertEquals("[cast , cast(], [distinct , distinct(], [drop , drop(], [select , select(], [extend , extend(], [filter , filter(], [from , from(], [groupBy , groupBy(], [pivot , pivot(], [join , join(], [limit , limit(], [rename , rename(], [size , size(], [slice , slice(], [sort , sort(]", checkResultNoException(new Completer("###Relational\nDatabase a::A(Table t(col VARCHAR(200)))", Lists.mutable.with(new RelationalCompleterExtension())).complete("#>{a::A.t}#->")));
Assert.assertEquals("[cast , cast(], [distinct , distinct(], [drop , drop(], [select , select(], [extend , extend(], [filter , filter(], [from , from(], [groupBy , groupBy(], [pivot , pivot(], [join , join(], [asOfJoin , asOfJoin(], [limit , limit(], [rename , rename(], [size , size(], [slice , slice(], [sort , sort(]", checkResultNoException(new Completer("###Relational\nDatabase a::A(Table t(col VARCHAR(200)))", Lists.mutable.with(new RelationalCompleterExtension())).complete("#>{a::A.t}#->")));
Assert.assertEquals("[select , select(], [size , size(], [slice , slice(], [sort , sort(]", checkResultNoException(new Completer("###Relational\nDatabase a::A(Table t(col VARCHAR(200)))", Lists.mutable.with(new RelationalCompleterExtension())).complete("#>{a::A.t}#->s")));
}

Expand Down Expand Up @@ -199,6 +199,24 @@ public void testJoin()
Assert.assertEquals("COMPILATION error at [6:14-17]: \"The relation contains duplicates: [val, col]\"", new Completer("###Relational\nDatabase a::A(Table t(col VARCHAR(200), val INT))", Lists.mutable.with(new RelationalCompleterExtension())).complete("#>{a::A.t}#->join(#>{a::A.t}#, JoinKind.INNER, {a,b|$a.val == $b.val})->").getEngineException().toPretty());
}

//------
// AsOfJoin
//------
@Test
public void testAsOfJoin()
{
Assert.assertEquals("[a::A , >{a::A]", checkResultNoException(new Completer("###Relational\nDatabase a::A(Table t(col VARCHAR(200), val INT))", Lists.mutable.with(new RelationalCompleterExtension())).complete("#>{a::A.t}#->asOfJoin(#>")));
Assert.assertEquals("[t , t}]", checkResultNoException(new Completer("###Relational\nDatabase a::A(Table t(col VARCHAR(200), val INT))", Lists.mutable.with(new RelationalCompleterExtension())).complete("#>{a::A.t}#->asOfJoin(#>{a::A.")));
Assert.assertEquals("[col , col], [val , val]", checkResultNoException(new Completer("###Relational\nDatabase a::A(Table t(col VARCHAR(200), val INT))", Lists.mutable.with(new RelationalCompleterExtension())).complete("#>{a::A.t}#->asOfJoin(#>{a::A.t}#, {a,b|$a.")));
Assert.assertEquals("", checkResultNoException(new Completer("###Relational\nDatabase a::A(Table t(col VARCHAR(200), val INT))", Lists.mutable.with(new RelationalCompleterExtension())).complete("#>{a::A.t}#->asOfJoin(#>{a::A.t}#,")));
Assert.assertEquals("[k , k], [o , o]", checkResultNoException(new Completer("###Relational\nDatabase a::A(Table t2(k VARCHAR(200), o INT) Table t(col VARCHAR(200), val INT))", Lists.mutable.with(new RelationalCompleterExtension())).complete("#>{a::A.t}#->asOfJoin(#>{a::A.t2}#, {a,b|$a.col == $b.")));
Assert.assertEquals("", checkResultNoException(new Completer("###Relational\nDatabase a::A(Table t2(k VARCHAR(200), o INT) Table t(col VARCHAR(200), val INT))", Lists.mutable.with(new RelationalCompleterExtension())).complete("#>{a::A.t}#->asOfJoin(#>{a::A.t2}#, {a,b|$a.col == $b.k}, ")));
Assert.assertEquals("[k , k], [o , o]", checkResultNoException(new Completer("###Relational\nDatabase a::A(Table t2(k VARCHAR(200), o INT) Table t(col VARCHAR(200), val INT))", Lists.mutable.with(new RelationalCompleterExtension())).complete("#>{a::A.t}#->asOfJoin(#>{a::A.t2}#, {a,b|$a.col == $b.k}, {b,c|$c.")));
Assert.assertEquals("[col , col], [val , val]", checkResultNoException(new Completer("###Relational\nDatabase a::A(Table t2(k VARCHAR(200), o INT) Table t(col VARCHAR(200), val INT))", Lists.mutable.with(new RelationalCompleterExtension())).complete("#>{a::A.t}#->asOfJoin(#>{a::A.t2}#, {a,b|$a.col == $b.k}, {b,c|$b.")));
Assert.assertEquals("[filter , filter(]", checkResultNoException(new Completer("###Relational\nDatabase a::A(Table t2(k VARCHAR(200), o INT) Table t(col VARCHAR(200), val INT))", Lists.mutable.with(new RelationalCompleterExtension())).complete("#>{a::A.t}#->asOfJoin(#>{a::A.t2}#, {a,b|$a.col == $b.k}, {b,c|$c.k == $b.col})->fil")));
Assert.assertEquals("COMPILATION error at [6:14-21]: \"The relation contains duplicates: [val, col]\"", new Completer("###Relational\nDatabase a::A(Table t(col VARCHAR(200), val INT))", Lists.mutable.with(new RelationalCompleterExtension())).complete("#>{a::A.t}#->asOfJoin(#>{a::A.t}#, {a,b|$a.val == $b.val})->").getEngineException().toPretty());
}

//--------
// Select
//--------
Expand Down
Loading

0 comments on commit fbcbc42

Please sign in to comment.