limit hints for query supported by delta-kernel

python/delta_sharing/tests/test_delta_sharing.py

@@ -354,10 +354,10 @@ def test_get_table_protocol(profile_path: str):
                 {
                     "eventTime": [
                         pd.Timestamp("2021-04-28 23:36:51.945"),
-                        pd.Timestamp("2021-04-28 23:36:47.599"),
+                        pd.Timestamp("2021-04-28 23:35:53.156"),
                     ],
                     "date": [date(2021, 4, 28), date(2021, 4, 28)],
-                    "type": ["bar", "foo"],
+                    "type": ["bar", None],
                 }
             ),
             id="limit 2",
@@ -370,11 +370,11 @@ def test_get_table_protocol(profile_path: str):
                 {
                     "eventTime": [
                         pd.Timestamp("2021-04-28 23:36:51.945"),
-                        pd.Timestamp("2021-04-28 23:36:47.599"),
                         pd.Timestamp("2021-04-28 23:35:53.156"),
+                        pd.Timestamp("2021-04-28 23:36:47.599")
                     ],
                     "date": [date(2021, 4, 28), date(2021, 4, 28), date(2021, 4, 28)],
-                    "type": ["bar", "foo", None],
+                    "type": ["bar", None, "foo"],
                 }
             ),
             id="limit 3",
@@ -387,11 +387,11 @@ def test_get_table_protocol(profile_path: str):
                 {
                     "eventTime": [
                         pd.Timestamp("2021-04-28 23:36:51.945"),
-                        pd.Timestamp("2021-04-28 23:36:47.599"),
                         pd.Timestamp("2021-04-28 23:35:53.156"),
+                        pd.Timestamp("2021-04-28 23:36:47.599")
                     ],
                     "date": [date(2021, 4, 28), date(2021, 4, 28), date(2021, 4, 28)],
-                    "type": ["bar", "foo", None],
+                    "type": ["bar", None, "foo"],
                 }
             ),
             id="limit 4",

server/src/main/scala/io/delta/kernelsharedtable/DeltaSharedTableKernel.scala

@@ -81,6 +81,9 @@ class DeltaSharedTableKernel(
   // from ColumnarBatch.
   private var scanFileFieldOrdinals: ScanFileFieldOrdinals = _
 
+  private var numRecords = 0L
+  private var earlyTermination = false
+
   private val fileSigner = withClassLoader {
     val tablePath = new Path(tableConfig.getLocation)
     val conf = new Configuration()
@@ -283,6 +286,12 @@ class DeltaSharedTableKernel(
       QueryTypes.QueryLatestSnapshot
     }
 
+    val globalLimitHint = if (predicateHints.isEmpty && jsonPredicateHints.isEmpty) {
+      limitHint
+    } else {
+      None
+    }
+
     val (table, engine) = getTableAndEngine()
 
     val snapshot = getSharedTableSnapshot(
@@ -309,12 +318,12 @@ class DeltaSharedTableKernel(
     )
 
     if (includeFiles) {
-      if (predicateHints.isEmpty && jsonPredicateHints.isEmpty && limitHint.isEmpty) {
+      if (predicateHints.isEmpty && jsonPredicateHints.isEmpty) {
         val scanBuilder = snapshot.kernelSnapshot.getScanBuilder(engine)
         val scan = scanBuilder.build()
         val scanFilesIter = scan.asInstanceOf[ScanImpl].getScanFiles(engine, true)
         try {
-          while (scanFilesIter.hasNext) {
+          while (scanFilesIter.hasNext && !earlyTermination) {
             val nextResponse = processBatchByColumnVector(
               scanFilesIter.next,
               snapshot.version,
@@ -323,7 +332,8 @@ class DeltaSharedTableKernel(
               table,
               engine,
               isVersionQuery,
-              snapshot
+              snapshot,
+              globalLimitHint
             )
             if (nextResponse.nonEmpty) {
               actions = actions ++ nextResponse
@@ -366,7 +376,8 @@ class DeltaSharedTableKernel(
       table: Table,
       engine: Engine,
       isVersionQuery: Boolean,
-      snapshot: SharedTableSnapshot): Seq[Object] = {
+      snapshot: SharedTableSnapshot,
+      limitHint: Option[Long]): Seq[Object] = {
 
     val batchData = scanFileBatch.getData
     val batchSize = batchData.getSize
@@ -376,6 +387,10 @@ class DeltaSharedTableKernel(
     val dataPath = new Path(table.getPath(engine))
 
     for (rowId <- 0 until batchSize) {
+      if (limitHint.exists(_ <= numRecords)) {
+        earlyTermination = true
+        return addFileObjects
+      }
       val isSelected = !selectionVector.isPresent ||
         (!selectionVector.get.isNullAt(rowId) && selectionVector.get.getBoolean(rowId))
       if (isSelected) {
@@ -636,6 +651,8 @@ class DeltaSharedTableKernel(
       if (addFileColumnVectors.stats.isNullAt(rowId)) null
       else addFileColumnVectors.stats.getString(rowId)
 
+    numRecords += JsonUtils.extractNumRecords(stats).getOrElse(0L)
+
     if (respondedFormat == DeltaSharedTableKernel.RESPONSE_FORMAT_DELTA) {
       DeltaResponseFileAction(
         // Using sha256 instead of m5 because it's less likely to have key collision.

server/src/main/scala/io/delta/sharing/server/DeltaSharingService.scala

@@ -499,7 +499,6 @@ class DeltaSharingService(serverConfig: ServerConfig) {
       val queryResult = if (
         request.predicateHints.isEmpty
           && request.jsonPredicateHints.isEmpty
-          && request.limitHint.isEmpty
           && request.includeRefreshToken.isEmpty
           && request.maxFiles.isEmpty
           && request.startingVersion.isEmpty

server/src/main/scala/io/delta/sharing/server/util/JsonUtils.scala

@@ -22,6 +22,8 @@ import com.fasterxml.jackson.annotation.JsonInclude.Include
 import com.fasterxml.jackson.databind.{DeserializationFeature, ObjectMapper}
 import com.fasterxml.jackson.module.scala.DefaultScalaModule
 import com.fasterxml.jackson.module.scala.experimental.ScalaObjectMapper
+import org.json4s.JsonAST.JInt
+import org.json4s.jackson.JsonMethods.parse
 
 object JsonUtils {
   /** Used to convert between classes and JSON. */
@@ -44,4 +46,18 @@ object JsonUtils {
   def fromJson[T: Manifest](json: String): T = {
     mapper.readValue[T](json)
   }
+
+  /** Parse the `numRecords` field from the stats json of the AddFile
+   * Returns 0 if the field is not found or if stats is empty. */
+  def extractNumRecords(stats: String): Option[Long] = {
+    if (stats != null && !stats.isEmpty()) {
+      val numRecordsField = parse(stats) \ "numRecords"
+      numRecordsField match {
+        case JInt(numRecords) => Some(numRecords.toLong)
+        case _ => None
+      }
+    } else {
+      None
+    }
+  }
 }

server/src/test/scala/io/delta/sharing/server/DeltaSharingServiceSuite.scala

@@ -1362,6 +1362,54 @@ class DeltaSharingServiceSuite extends FunSuite with BeforeAndAfterAll {
     verifyPreSignedUrl(actualFiles(2).url, 778)
   }
 
+  integrationTest("table3 - different data file schemas limit hint - /shares/{share}/schemas/{schema}/tables/{table}/query") {
+    val p =
+      """
+        |{
+        | "limitHint": 2
+        |}
+        |""".stripMargin
+
+    val response = readNDJson(requestPath("/shares/share1/schemas/default/tables/table3/query"), Some("POST"), Some(p), Some(4))
+    print(response)
+    val lines = response.split("\n")
+    val protocol = lines(0)
+    val metadata = lines(1)
+    val expectedProtocol = Protocol(minReaderVersion = 1).wrap
+    assert(expectedProtocol == JsonUtils.fromJson[SingleAction](protocol))
+    val expectedMetadata = Metadata(
+      id = "7ba6d727-a578-4234-a138-953f790b427c",
+      format = Format(),
+      schemaString = """{"type":"struct","fields":[{"name":"eventTime","type":"timestamp","nullable":true,"metadata":{}},{"name":"date","type":"date","nullable":true,"metadata":{}},{"name":"type","type":"string","nullable":true,"metadata":{}}]}""",
+      partitionColumns = Seq("date")).wrap
+    assert(expectedMetadata == JsonUtils.fromJson[SingleAction](metadata))
+    val files = lines.drop(2)
+    val actualFiles = files.map(f => JsonUtils.fromJson[SingleAction](f).file)
+    assert(actualFiles.size == 2)
+    val expectedFiles = Seq(
+      AddFile(
+        url = actualFiles(0).url,
+        expirationTimestamp = actualFiles(0).expirationTimestamp,
+        id = "db213271abffec6fd6c7fc2aad9d4b3f",
+        partitionValues = Map("date" -> "2021-04-28"),
+        size = 778,
+        stats = """{"numRecords":1,"minValues":{"eventTime":"2021-04-28T23:36:51.945Z","type":"bar"},"maxValues":{"eventTime":"2021-04-28T23:36:51.945Z","type":"bar"},"nullCount":{"eventTime":0,"type":0}}"""
+      ),
+      AddFile(
+        url = actualFiles(1).url,
+        expirationTimestamp = actualFiles(1).expirationTimestamp,
+        id = "a892a55d770ee70b34ffb2ebf7dc2fd0",
+        partitionValues = Map("date" -> "2021-04-28"),
+        size = 573,
+        stats = """{"numRecords":1,"minValues":{"eventTime":"2021-04-28T23:35:53.156Z"},"maxValues":{"eventTime":"2021-04-28T23:35:53.156Z"},"nullCount":{"eventTime":0}}"""
+      )
+    )
+    assert(actualFiles.count(_.expirationTimestamp != null) == 2)
+    assert(expectedFiles == actualFiles.toList)
+    verifyPreSignedUrl(actualFiles(0).url, 778)
+    verifyPreSignedUrl(actualFiles(1).url, 573)
+  }
+
   integrationTest("case insensitive") {
     val response = readNDJson(requestPath("/shares/sHare1/schemas/deFault/tables/taBle3/metadata"), expectedTableVersion = Some(4))
     val Array(protocol, metadata) = response.split("\n")