From 2713f084569f4c720ec4b954f154accb9d560938 Mon Sep 17 00:00:00 2001
From: Sean Kao <seankao@amazon.com>
Date: Mon, 16 Oct 2023 15:19:34 -0700
Subject: [PATCH] query rewrite for nexus skipping index

Signed-off-by: Sean Kao <seankao@amazon.com>
---
 build.sbt                                     |   4 ++
 .../lib/AWSLogsDALConnectorSpark-1.0.jar      | Bin 0 -> 2193 bytes
 .../ApplyFlintSparkSkippingIndex.scala        |  54 +++++++++++++++++-
 3 files changed, 57 insertions(+), 1 deletion(-)
 create mode 100644 flint-spark-integration/lib/AWSLogsDALConnectorSpark-1.0.jar

diff --git a/build.sbt b/build.sbt
index 0dcfb8af7..4cb825f1a 100644
--- a/build.sbt
+++ b/build.sbt
@@ -139,6 +139,10 @@ lazy val flintSparkIntegration = (project in file("flint-spark-integration"))
         val oldStrategy = (assembly / assemblyMergeStrategy).value
         oldStrategy(x)
     },
+    assembly / assemblyExcludedJars := {
+      val cp = (assembly / fullClasspath).value
+      cp filter { file => file.data.getName.contains("AWSLogsDALConnectorSpark")}
+    },
     assembly / test := (Test / test).value)
 
 // Test assembly package with integration test.
diff --git a/flint-spark-integration/lib/AWSLogsDALConnectorSpark-1.0.jar b/flint-spark-integration/lib/AWSLogsDALConnectorSpark-1.0.jar
new file mode 100644
index 0000000000000000000000000000000000000000..30c3fd3c40a7c71458090c4751bea6b8715f78c3
GIT binary patch
literal 2193
zcmWIWW@h1HVBp|j;ONQ?X8-~w5CH_7K<w)p;;8HC=cdoh5CB)0&=nuPsh;VbG*E#g
z5Q`uy^mX*}baM?3(erit?0e?4kGHPgMP6@Rt#fD2Zw@lJV*KFgqo+(rhN0Wa0yHc+
zKNrP9bT!;iHHo>2Rrz^%l`5esO)M|Y$xkm%Ni0cBOioTMF3w4;OfAA|f-X@eB<JVl
zr6!l;7Zu|*%$QWeiVG5pvN4Pdd$i(+2r#64fg#~Uo{>JlU<gS}%1PBr&PgmT-WnQP
zeA`9DHrsmQ(p{djrUf`Vdn)v<t>L&dg~6=R#Kh!?0NeC44)d>l>RTK(Eygp?aDVgv
z3+Zx3i3x1?DjrOfXSoxr-S^h2SB*#i(%rki3*YT5fB*02cm4g0HxA5e6uBg2WZ9EA
zC-IeF{VCn))x6pJzj@pcJX~?}<D87iim@7UFHSgJPk2(mY1y{vt-0>KC-y()9i890
z?mowBn@4xHxk~KXqs(;rn{EG}AC4KVx^gz#w3d8nxE|xU&iQwMiRIy>qx@20`Grx(
zvmOdx@3PVrjJm3kar+|kByXcXE_{<`BwhFu_T)$034_SM6}Kloj*9;fcJhG9x~v`3
z4r|XyG;rOU+qq1>WQ|%s-`fkT3V#=#=2v=sB~P29qP?c~^M*$)Q@%MpOiZ286SPr(
zQb=&ay5|#)XY;Tv-T!z;(ssYPWp0`GP9%mo&Cg-om*R1yT;}05=jL;Hr?~c<{IFrd
zgzkk)eqWq4?XhW_X&1ZHv(DGSd-L~fxI4+E<5!(GOU$Rr_tnuFNxWxHo{LD>VJj4{
ze#10h)#Z2AW?hf`vGqYI+um0uHga2;wZ6=C)Ofmi!?*W2(<JW6B^9<T3^|?iyHxS_
zfsaR~zfJRHDQxx)ubI7@-KOA))x9$=?R=pbn?J?Rn!YGb*3t6tnFNnjAu`{JOc!li
z{N`%ZVX3bXi@dTU)L%0)nXWOi*&nf7MdkmUI<FVYD-IO=Z?gUC#BBWM$9}cihdf`~
zI#zu&&=%|BG`w_1VfWw3mSsx=WKPXUOs>%Xcz34S^?2TMl39zLXXV!>ccuL*=MkP)
z5w>+_`S!OvN(Bu}&Lu=P-gvg~##g6@t$t5zMDI>mWw*H3kKNUw`QR&V7me%sLLKim
z9G_jXeBIm)#ue<vW|nI#^=5X8SFG*o2>I6b{nYbD`<X?re&x+)d7-hcch!Bq9K}pk
zvufw0P`3q!vzEF~Ym+qkVsLS{r~dmbHHS0ai=EWexio9ugnPRjUf*miZceF^^U%9=
zF8+)B;`)m=Y{t|37R)^DzGKfPcC|B8=NCTZSa0WkE#-^(=b6fT{t0WQCtNC@cYxt~
z=G1mSMfCz_*4amWt!6L2@Hs-M_iG-5PBLr2B$jhiR2D9uGUd3|C9g>Z#v455b)Hw_
zEYq@!Ub4B|<DBwWe%<3HDW-c?Kit-4{blNdT{HJB|8%G0bB$9`+rM2qx10&<``C6y
zdF>~IDeqZNJ)3qYyG*pY;;UTT-uCctjmxzfU*)`F7y2AGxTk)a|AS(!YL98Cox*SJ
zXZviAMeo;^?M#Y!JDbbVJj%@NK!B5_-(#EKR|UVt%E({5mzHFot^f72p`GW9^T$rq
zY%4I|HoYQ9f1|$kZ@t4e<6Ji#`D!-TxTdR`b7D@fOOa+_YMI-u_YXF=^t_vU@<r{f
z+o6}9|7QyDW@Hj!#$7o83mza4V0h~YqTy8qvNmj01xPUnENRpPl5lNE6$i2j=w&~s
zet-aw#brPysJ6hT15{f;07%CyAQPkmT6f^n2&y|E0HpCKkcp%bM@@pyJWx#n0U-0(
z8Hq6usb0ZnG^k!d7#%{Y(U4jO-2(Jl24=;Q#^vN$fL7-Oc(byB%mg+)7!rW34rwqC
E0F3Psr2qf`

literal 0
HcmV?d00001

diff --git a/flint-spark-integration/src/main/scala/org/opensearch/flint/spark/skipping/ApplyFlintSparkSkippingIndex.scala b/flint-spark-integration/src/main/scala/org/opensearch/flint/spark/skipping/ApplyFlintSparkSkippingIndex.scala
index 11f8ad304..5a2f97038 100644
--- a/flint-spark-integration/src/main/scala/org/opensearch/flint/spark/skipping/ApplyFlintSparkSkippingIndex.scala
+++ b/flint-spark-integration/src/main/scala/org/opensearch/flint/spark/skipping/ApplyFlintSparkSkippingIndex.scala
@@ -5,14 +5,19 @@
 
 package org.opensearch.flint.spark.skipping
 
+import com.amazon.awslogsdataaccesslayer.connectors.spark.LogsTable
 import org.opensearch.flint.spark.FlintSpark
-import org.opensearch.flint.spark.skipping.FlintSparkSkippingIndex.{getSkippingIndexName, SKIPPING_INDEX_TYPE}
+import org.opensearch.flint.spark.skipping.FlintSparkSkippingIndex.{getSkippingIndexName, FILE_PATH_COLUMN, SKIPPING_INDEX_TYPE}
 
+import org.apache.spark.sql.Column
 import org.apache.spark.sql.catalyst.catalog.CatalogTable
 import org.apache.spark.sql.catalyst.expressions.{And, Expression, Or, Predicate}
 import org.apache.spark.sql.catalyst.plans.logical.{Filter, LogicalPlan}
 import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.connector.catalog.{CatalogPlugin, Identifier}
 import org.apache.spark.sql.execution.datasources.{HadoopFsRelation, LogicalRelation}
+import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation
+import org.apache.spark.sql.flint.config.FlintSparkConf
 import org.apache.spark.sql.flint.qualifyTableName
 
 /**
@@ -57,6 +62,48 @@ class ApplyFlintSparkSkippingIndex(flint: FlintSpark) extends Rule[LogicalPlan]
       } else {
         filter
       }
+    case filter @ Filter(
+          condition: Predicate,
+          relation @ DataSourceV2Relation(table, _, Some(catalog), Some(identifier), _))
+        if hasNoDisjunction(condition) &&
+          // Check if query plan already rewritten
+          table.isInstanceOf[LogsTable] && !table.asInstanceOf[LogsTable].hasLogFileIds() =>
+      val index = flint.describeIndex(getIndexName(catalog, identifier))
+      if (index.exists(_.kind == SKIPPING_INDEX_TYPE)) {
+        val skippingIndex = index.get.asInstanceOf[FlintSparkSkippingIndex]
+        val indexFilter = rewriteToIndexFilter(skippingIndex, condition)
+        /*
+         * Replace original LogsTable with a new one with log file idx:
+         *  Filter(a=b)
+         *  |- DataSourceV2Relation(A)
+         *     |- LogsTable <== replaced with a new LogsTable with log file ids
+         */
+        if (indexFilter.isDefined) {
+          val indexScan = flint.queryIndex(skippingIndex.name())
+          val selectedFiles =
+            // Non hybrid scan
+            // TODO: refactor common logic with file-based skipping index
+            indexScan
+              .filter(new Column(indexFilter.get))
+              .select(FILE_PATH_COLUMN)
+              .collect
+              .map(_.getString(0))
+
+          // Construct LogsTable with selectedFiles as its log file ids
+          // It will build scan operator using these log file ids
+          val logsTable = table.asInstanceOf[LogsTable]
+          val newTable = new LogsTable(
+            logsTable.getSchema(),
+            logsTable.getOptions(),
+            selectedFiles,
+            logsTable.getProcessedFields())
+          filter.copy(child = relation.copy(table = newTable))
+        } else {
+          filter
+        }
+      } else {
+        filter
+      }
   }
 
   private def getIndexName(table: CatalogTable): String = {
@@ -67,6 +114,11 @@ class ApplyFlintSparkSkippingIndex(flint: FlintSpark) extends Rule[LogicalPlan]
     getSkippingIndexName(qualifiedTableName)
   }
 
+  private def getIndexName(catalog: CatalogPlugin, identifier: Identifier): String = {
+    val qualifiedTableName = s"${catalog.name}.${identifier}"
+    getSkippingIndexName(qualifiedTableName)
+  }
+
   private def hasNoDisjunction(condition: Expression): Boolean = {
     condition.collectFirst { case Or(_, _) =>
       true