diff --git a/flint-spark-integration/src/main/scala/org/opensearch/flint/spark/FlintSparkIndexBuilder.scala b/flint-spark-integration/src/main/scala/org/opensearch/flint/spark/FlintSparkIndexBuilder.scala index 7171bf1af..95e351f7d 100644 --- a/flint-spark-integration/src/main/scala/org/opensearch/flint/spark/FlintSparkIndexBuilder.scala +++ b/flint-spark-integration/src/main/scala/org/opensearch/flint/spark/FlintSparkIndexBuilder.scala @@ -34,6 +34,14 @@ abstract class FlintSparkIndexBuilder(flint: FlintSpark) { .toMap } + /** + * Add index options. + * + * @param options + * index options + * @return + * builder + */ def options(options: FlintSparkIndexOptions): this.type = { this.indexOptions = options this diff --git a/flint-spark-integration/src/main/scala/org/opensearch/flint/spark/skipping/FlintSparkSkippingIndex.scala b/flint-spark-integration/src/main/scala/org/opensearch/flint/spark/skipping/FlintSparkSkippingIndex.scala index c871a1900..2eebffebf 100644 --- a/flint-spark-integration/src/main/scala/org/opensearch/flint/spark/skipping/FlintSparkSkippingIndex.scala +++ b/flint-spark-integration/src/main/scala/org/opensearch/flint/spark/skipping/FlintSparkSkippingIndex.scala @@ -33,9 +33,9 @@ import org.apache.spark.sql.types.StructType * @param indexedColumns * indexed column list */ -case class FlintSparkSkippingIndex( +class FlintSparkSkippingIndex( tableName: String, - indexedColumns: Seq[FlintSparkSkippingStrategy], + val indexedColumns: Seq[FlintSparkSkippingStrategy], override val options: FlintSparkIndexOptions = empty) extends FlintSparkIndex { diff --git a/flint-spark-integration/src/main/scala/org/opensearch/flint/spark/sql/FlintSparkSqlAstBuilder.scala b/flint-spark-integration/src/main/scala/org/opensearch/flint/spark/sql/FlintSparkSqlAstBuilder.scala index 3a67625de..9f3539b09 100644 --- a/flint-spark-integration/src/main/scala/org/opensearch/flint/spark/sql/FlintSparkSqlAstBuilder.scala +++ b/flint-spark-integration/src/main/scala/org/opensearch/flint/spark/sql/FlintSparkSqlAstBuilder.scala @@ -5,11 +5,13 @@ package org.opensearch.flint.spark.sql -import org.antlr.v4.runtime.tree.RuleNode +import org.antlr.v4.runtime.tree.{ParseTree, RuleNode} import org.opensearch.flint.spark.FlintSpark import org.opensearch.flint.spark.sql.covering.FlintSparkCoveringIndexAstBuilder import org.opensearch.flint.spark.sql.skipping.FlintSparkSkippingIndexAstBuilder +import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan + /** * Flint Spark AST builder that builds Spark command for Flint index statement. This class mix-in * all other AST builders and provides util methods. @@ -20,6 +22,10 @@ class FlintSparkSqlAstBuilder with FlintSparkCoveringIndexAstBuilder with SparkSqlAstBuilder { + override def visit(tree: ParseTree): LogicalPlan = { + tree.accept(this).asInstanceOf[LogicalPlan] + } + override def aggregateResult(aggregate: AnyRef, nextResult: AnyRef): AnyRef = if (nextResult != null) nextResult else aggregate } diff --git a/flint-spark-integration/src/main/scala/org/opensearch/flint/spark/sql/FlintSparkSqlParser.scala b/flint-spark-integration/src/main/scala/org/opensearch/flint/spark/sql/FlintSparkSqlParser.scala index 4bc828644..0fa146b9d 100644 --- a/flint-spark-integration/src/main/scala/org/opensearch/flint/spark/sql/FlintSparkSqlParser.scala +++ b/flint-spark-integration/src/main/scala/org/opensearch/flint/spark/sql/FlintSparkSqlParser.scala @@ -54,7 +54,7 @@ class FlintSparkSqlParser(sparkParser: ParserInterface) extends ParserInterface override def parsePlan(sqlText: String): LogicalPlan = parse(sqlText) { flintParser => try { - flintAstBuilder.visit(flintParser.singleStatement()).asInstanceOf[LogicalPlan] + flintAstBuilder.visit(flintParser.singleStatement()) } catch { // Fall back to Spark parse plan logic if flint cannot parse case _: ParseException => sparkParser.parsePlan(sqlText) diff --git a/flint-spark-integration/src/main/scala/org/opensearch/flint/spark/sql/SparkSqlAstBuilder.scala b/flint-spark-integration/src/main/scala/org/opensearch/flint/spark/sql/SparkSqlAstBuilder.scala index 8a5855d47..d4e4db104 100644 --- a/flint-spark-integration/src/main/scala/org/opensearch/flint/spark/sql/SparkSqlAstBuilder.scala +++ b/flint-spark-integration/src/main/scala/org/opensearch/flint/spark/sql/SparkSqlAstBuilder.scala @@ -15,7 +15,8 @@ import org.opensearch.flint.spark.sql.FlintSparkSqlExtensionsParser.{PropertyKey import org.apache.spark.sql.catalyst.parser.ParserUtils.string /** - * AST builder that builds for common rule in Spark SQL grammar. + * AST builder that builds for common rule in Spark SQL grammar. The main logic is modified slightly + * from Spark AstBuilder code. */ trait SparkSqlAstBuilder extends FlintSparkSqlExtensionsVisitor[AnyRef] {