-
Notifications
You must be signed in to change notification settings - Fork 34
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Signed-off-by: Kacper Trochimiak <kacper.trochimiak@eliatra.com>
- Loading branch information
1 parent
2e16414
commit 714bfcf
Showing
8 changed files
with
224 additions
and
2 deletions.
There are no files selected for viewing
62 changes: 62 additions & 0 deletions
62
.../src/integration/scala/org/opensearch/flint/spark/ppl/FlintSparkPPLTrendlineITSuite.scala
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,62 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| package org.opensearch.flint.spark.ppl | ||
|
|
||
| import org.apache.spark.sql.catalyst.analysis.UnresolvedRelation | ||
| import org.apache.spark.sql.catalyst.plans.logical._ | ||
| import org.apache.spark.sql.streaming.StreamTest | ||
| import org.apache.spark.sql.{QueryTest, Row} | ||
|
|
||
| class FlintSparkPPLTrendlineITSuite | ||
| extends QueryTest | ||
| with LogicalPlanTestUtils | ||
| with FlintPPLSuite | ||
| with StreamTest { | ||
|
|
||
| /** Test table and index name */ | ||
| private val testTable = "spark_catalog.default.flint_ppl_test" | ||
|
|
||
| override def beforeAll(): Unit = { | ||
| super.beforeAll() | ||
|
|
||
| // Create test table | ||
| createPartitionedStateCountryTable(testTable) | ||
| } | ||
|
|
||
| protected override def afterEach(): Unit = { | ||
| super.afterEach() | ||
| // Stop all streaming jobs if any | ||
| spark.streams.active.foreach { job => | ||
| job.stop() | ||
| job.awaitTermination() | ||
| } | ||
| } | ||
|
|
||
| test("trendline sma") { | ||
| val frame = sql(s""" | ||
| | source = $testTable | trendline sma(2, age) as first_age_sma sma(3, age) as second_age_sma | fields name, first_age_sma, second_age_sma | ||
| | """.stripMargin) | ||
|
|
||
| // Retrieve the results | ||
| val results: Array[Row] = frame.collect() | ||
| // Define the expected results | ||
| val expectedResults: Array[Row] = Array() | ||
|
|
||
| // Convert actual results to a Set for quick lookup | ||
| val resultsSet: Set[Row] = results.toSet | ||
| // Check that each expected row is present in the actual results | ||
| expectedResults.foreach { expectedRow => | ||
| assert(resultsSet.contains(expectedRow), s"Expected row $expectedRow not found in results") | ||
| } | ||
| // Retrieve the logical plan | ||
| val logicalPlan: LogicalPlan = | ||
| frame.queryExecution.commandExecuted.asInstanceOf[CommandResult].commandLogicalPlan | ||
| // Define the expected logical plan | ||
| val expectedPlan: LogicalPlan = Project(Seq(), new UnresolvedRelation(Seq("spark_catalog", "default", "flint_ppl_test"))) | ||
| // Compare the two plans | ||
| comparePlans(logicalPlan, expectedPlan, checkAnalysis = false) | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
65 changes: 65 additions & 0 deletions
65
ppl-spark-integration/src/main/java/org/opensearch/sql/ast/tree/Trendline.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,65 @@ | ||
| package org.opensearch.sql.ast.tree; | ||
|
|
||
| import com.google.common.collect.ImmutableList; | ||
| import lombok.EqualsAndHashCode; | ||
| import lombok.Getter; | ||
| import lombok.RequiredArgsConstructor; | ||
| import lombok.ToString; | ||
| import org.opensearch.sql.ast.AbstractNodeVisitor; | ||
| import org.opensearch.sql.ast.Node; | ||
| import org.opensearch.sql.ast.expression.UnresolvedExpression; | ||
|
|
||
| import java.util.List; | ||
|
|
||
| @ToString | ||
| @Getter | ||
| @RequiredArgsConstructor | ||
| @EqualsAndHashCode(callSuper = false) | ||
| public class Trendline extends UnresolvedPlan { | ||
|
|
||
| private UnresolvedPlan child; | ||
| private final List<UnresolvedExpression> computations; | ||
|
|
||
| @Override | ||
| public UnresolvedPlan attach(UnresolvedPlan child) { | ||
| this.child = child; | ||
| return this; | ||
| } | ||
|
|
||
| @Override | ||
| public List<? extends Node> getChild() { | ||
| return ImmutableList.of(child); | ||
| } | ||
|
|
||
| @Override | ||
| public <T, C> T accept(AbstractNodeVisitor<T, C> visitor, C context) { | ||
| return visitor.visitTrendline(this, context); | ||
| } | ||
|
|
||
| @Getter | ||
| public static class TrendlineComputation extends UnresolvedExpression { | ||
|
|
||
| private final Integer numberOfDataPoints; | ||
| private final UnresolvedExpression dataField; | ||
| private final String alias; | ||
| private final TrendlineType computationType; | ||
|
|
||
| public TrendlineComputation(Integer numberOfDataPoints, UnresolvedExpression dataField, String alias, String computationType) { | ||
| this.numberOfDataPoints = numberOfDataPoints; | ||
| this.dataField = dataField; | ||
| this.alias = alias; | ||
| this.computationType = Trendline.TrendlineType.valueOf(computationType.toUpperCase()); | ||
| } | ||
|
|
||
| @Override | ||
| public <R, C> R accept(AbstractNodeVisitor<R, C> nodeVisitor, C context) { | ||
| return nodeVisitor.visitTrendlineComputation(this, context); | ||
| } | ||
|
|
||
| } | ||
|
|
||
| public enum TrendlineType { | ||
| SMA, | ||
| WMA | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters