fixed pushFilters for ReadMode.Query (#37)

Author: rashtao

arangodb-spark-datasource-2.4/src/main/scala/org/apache/spark/sql/arangodb/datasource/reader/ArangoDataSourceReader.scala

@@ -39,6 +39,13 @@ class ArangoDataSourceReader(tableSchema: StructType, options: ArangoDBConf) ext
       .map(it => ArangoPartition.ofCollection(it._1, it._2, new PushDownCtx(readSchema(), appliedPushableFilters), options))
 
   override def pushFilters(filters: Array[Filter]): Array[Filter] = {
+    options.readOptions.readMode match {
+      case ReadMode.Collection => pushFiltersReadModeCollection(filters)
+      case ReadMode.Query => filters
+    }
+  }
+
+  private def pushFiltersReadModeCollection(filters: Array[Filter]): Array[Filter] = {
     // filters related to columnNameOfCorruptRecord are not pushed down
     val isCorruptRecordFilter = (f: Filter) => f.references.contains(options.readOptions.columnNameOfCorruptRecord)
     val ignoredFilters = filters.filter(isCorruptRecordFilter)

arangodb-spark-datasource-3.1/src/main/scala/org/apache/spark/sql/arangodb/datasource/reader/ArangoScanBuilder.scala

@@ -1,7 +1,7 @@
 package org.apache.spark.sql.arangodb.datasource.reader
 
 import org.apache.spark.internal.Logging
-import org.apache.spark.sql.arangodb.commons.ArangoDBConf
+import org.apache.spark.sql.arangodb.commons.{ArangoDBConf, ReadMode}
 import org.apache.spark.sql.arangodb.commons.filter.{FilterSupport, PushableFilter}
 import org.apache.spark.sql.arangodb.commons.utils.PushDownCtx
 import org.apache.spark.sql.connector.read.{Scan, ScanBuilder, SupportsPushDownFilters, SupportsPushDownRequiredColumns}
@@ -22,6 +22,13 @@ class ArangoScanBuilder(options: ArangoDBConf, tableSchema: StructType) extends
   override def build(): Scan = new ArangoScan(new PushDownCtx(readSchema, appliedPushableFilters), options)
 
   override def pushFilters(filters: Array[Filter]): Array[Filter] = {
+    options.readOptions.readMode match {
+      case ReadMode.Collection => pushFiltersReadModeCollection(filters)
+      case ReadMode.Query => filters
+    }
+  }
+
+  private def pushFiltersReadModeCollection(filters: Array[Filter]): Array[Filter] = {
     // filters related to columnNameOfCorruptRecord are not pushed down
     val isCorruptRecordFilter = (f: Filter) => f.references.contains(options.readOptions.columnNameOfCorruptRecord)
     val ignoredFilters = filters.filter(isCorruptRecordFilter)

arangodb-spark-datasource-3.2/src/main/scala/org/apache/spark/sql/arangodb/datasource/reader/ArangoScanBuilder.scala

@@ -1,7 +1,7 @@
 package org.apache.spark.sql.arangodb.datasource.reader
 
 import org.apache.spark.internal.Logging
-import org.apache.spark.sql.arangodb.commons.ArangoDBConf
+import org.apache.spark.sql.arangodb.commons.{ArangoDBConf, ReadMode}
 import org.apache.spark.sql.arangodb.commons.filter.{FilterSupport, PushableFilter}
 import org.apache.spark.sql.arangodb.commons.utils.PushDownCtx
 import org.apache.spark.sql.connector.read.{Scan, ScanBuilder, SupportsPushDownFilters, SupportsPushDownRequiredColumns}
@@ -22,6 +22,13 @@ class ArangoScanBuilder(options: ArangoDBConf, tableSchema: StructType) extends
   override def build(): Scan = new ArangoScan(new PushDownCtx(readSchema, appliedPushableFilters), options)
 
   override def pushFilters(filters: Array[Filter]): Array[Filter] = {
+    options.readOptions.readMode match {
+      case ReadMode.Collection => pushFiltersReadModeCollection(filters)
+      case ReadMode.Query => filters
+    }
+  }
+
+  private def pushFiltersReadModeCollection(filters: Array[Filter]): Array[Filter] = {
     // filters related to columnNameOfCorruptRecord are not pushed down
     val isCorruptRecordFilter = (f: Filter) => f.references.contains(options.readOptions.columnNameOfCorruptRecord)
     val ignoredFilters = filters.filter(isCorruptRecordFilter)

integration-tests/src/test/scala/org/apache/spark/sql/arangodb/datasource/ReadTest.scala

@@ -5,7 +5,7 @@ import org.apache.spark.SparkException
 import org.apache.spark.sql.DataFrame
 import org.apache.spark.sql.arangodb.commons.ArangoDBConf
 import org.apache.spark.sql.functions.col
-import org.apache.spark.sql.types.{NullType, NumericType, StringType, StructField, StructType}
+import org.apache.spark.sql.types._
 import org.assertj.core.api.Assertions.{assertThat, catchThrowable}
 import org.assertj.core.api.ThrowableAssert.ThrowingCallable
 import org.junit.jupiter.api.Assumptions.assumeTrue
@@ -167,6 +167,28 @@ class ReadTest extends BaseSparkTest {
     assertThat(df.count()).isEqualTo(10)
   }
 
+  @ParameterizedTest
+  @MethodSource(Array("provideProtocolAndContentType"))
+  def readQueryWithFilter(protocol: String, contentType: String): Unit = {
+    val query =
+      """
+        |FOR i IN 1..10
+        | RETURN { idx: i, value: SHA1(i) }
+        |""".stripMargin.replaceAll("\n", "")
+
+    val df = spark.read
+      .format(BaseSparkTest.arangoDatasource)
+      .options(options + (
+        ArangoDBConf.QUERY -> query,
+        ArangoDBConf.PROTOCOL -> protocol,
+        ArangoDBConf.CONTENT_TYPE -> contentType
+      ))
+      .load()
+      .filter(col("idx") === 3)
+
+    assertThat(df.count()).isEqualTo(1)
+  }
+
   @ParameterizedTest
   @MethodSource(Array("provideProtocolAndContentType"))
   def reatTimeout(protocol: String, contentType: String): Unit = {