[SPARK-54305][SQL][PYTHON] Add admission control support to Python DataSource streaming API

This PR adds admission control support to Python streaming data sources by:

1. Enhanced Python API: Modified DataSourceStreamReader.latestOffset() to accept optional start and limit parameters
2. Scala Bridge: Updated PythonMicroBatchStream to validate admission control options and throw IllegalArgumentException for invalid values
3. Python-JVM Communication: Added function IDs and handlers in PythonStreamingSourceRunner for admission control parameters
4. Python Worker: Implemented latest_offset_with_report_func() with backward compatibility fallback
5. Documentation: Added comprehensive guide in python_data_source.rst
6. Tests: Added validation tests for invalid admission control parameters
7. Example: Created structured_blockchain_admission_control.py demonstrating the Full API

Key Benefits:
- Predictable performance with controlled batch sizes
- Rate limiting and backpressure support
- Feature parity with Scala DataSource capabilities
- Full backward compatibility (all parameters optional)

Author: jiteshsoni

examples/src/main/python/sql/streaming/structured_blockchain_admission_control.py

@@ -0,0 +1,167 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""
+Demonstrates admission control in Python streaming data sources.
+
+This example implements a simple blockchain-like streaming source that generates
+sequential blocks and shows how to use admission control to limit batch sizes.
+
+Usage: structured_blockchain_admission_control.py [<max-blocks-per-batch>]
+  <max-blocks-per-batch> Maximum number of blocks to process per microbatch (default: 10)
+
+Run the example:
+   `$ bin/spark-submit examples/src/main/python/sql/streaming/\\
+structured_blockchain_admission_control.py 5`
+
+The example will process blocks in controlled batches of 5, demonstrating admission control.
+"""
+import sys
+import time
+
+from pyspark.sql import SparkSession
+from pyspark.sql.datasource import DataSource, DataSourceStreamReader, InputPartition
+
+
+class SimpleBlockchainReader(DataSourceStreamReader):
+    """A simple streaming source that generates sequential blockchain blocks."""
+
+    def __init__(self, max_block=1000):
+        self.max_block = max_block
+        self.current_block = 0
+
+    def initialOffset(self):
+        """Start from block 0."""
+        return {"block": self.current_block}
+
+    def latestOffset(self, start=None, limit=None):
+        """
+        Return the latest offset, respecting admission control limits.
+
+        This demonstrates the key admission control pattern:
+        - Without limit: process all available blocks
+        - With maxRows limit: cap the end block to respect batch size
+        """
+        # Determine where we are now
+        if start is None:
+            start_block = self.current_block
+        else:
+            start_block = start["block"]
+
+        # Simulate blockchain growth - advance by 20 blocks each time
+        latest_available = min(start_block + 20, self.max_block)
+
+        # Apply admission control if configured
+        if limit and limit.get("type") == "maxRows":
+            max_blocks = limit["maxRows"]
+            # Cap at the configured limit
+            end_block = min(start_block + max_blocks, latest_available)
+            print(
+                f"  [Admission Control] Start: {start_block}, Available: {latest_available}, "
+                f"Capped: {end_block} (limit: {max_blocks})"
+            )
+            # Return tuple: (capped_offset, true_latest_offset)
+            return ({"block": end_block}, {"block": latest_available})
+        else:
+            # No limit - process all available
+            end_block = latest_available
+            print(f"  [No Limit] Start: {start_block}, End: {end_block}")
+            return {"block": end_block}
+
+    def partitions(self, start, end):
+        """Create a single partition for the block range."""
+        start_block = start["block"]
+        end_block = end["block"]
+        return [InputPartition(f"{start_block}:{end_block}".encode())]
+
+    def read(self, partition):
+        """Generate block data for the partition."""
+        # Parse the block range
+        range_str = partition.value.decode()
+        start_block, end_block = map(int, range_str.split(":"))
+
+        # Generate block data
+        for block_num in range(start_block, end_block):
+            # Simulate block data: block number, timestamp, simple hash
+            yield (
+                block_num,
+                int(time.time() * 1000),
+                f"0x{'0' * 60}{block_num:04x}",
+            )
+
+    def commit(self, end):
+        """Mark this offset as committed."""
+        pass
+
+
+class SimpleBlockchainSource(DataSource):
+    """Data source for simple blockchain streaming."""
+
+    @classmethod
+    def name(cls):
+        return "simple_blockchain"
+
+    def schema(self):
+        return "block_number INT, timestamp LONG, block_hash STRING"
+
+    def streamReader(self, schema):
+        return SimpleBlockchainReader(max_block=1000)
+
+
+if __name__ == "__main__":
+    max_blocks_per_batch = int(sys.argv[1]) if len(sys.argv) > 1 else 10
+
+    print(
+        f"""
+=================================================================
+Blockchain Streaming with Admission Control
+=================================================================
+Configuration:
+  - Max blocks per batch: {max_blocks_per_batch}
+  - Total blocks to generate: 1000
+
+Watch how admission control limits each microbatch to process
+only {max_blocks_per_batch} blocks at a time, even when more data is available.
+=================================================================
+"""
+    )
+
+    spark = (
+        SparkSession.builder.appName("StructuredBlockchainAdmissionControl").getOrCreate()
+    )
+
+    # Register the custom data source
+    spark.dataSource.register(SimpleBlockchainSource)
+
+    # Create streaming DataFrame with admission control
+    blocks = (
+        spark.readStream.format("simple_blockchain")
+        .option("maxRecordsPerBatch", str(max_blocks_per_batch))
+        .load()
+    )
+
+    # Show block statistics per microbatch
+    query = (
+        blocks.writeStream.outputMode("append")
+        .format("console")
+        .option("numRows", "20")
+        .option("truncate", "false")
+        .trigger(processingTime="3 seconds")
+        .start()
+    )
+
+    query.awaitTermination()

python/docs/source/tutorial/sql/python_data_source.rst

@@ -487,6 +487,101 @@ We can also use the same data source in streaming reader and writer
 
     query = spark.readStream.format("fake").load().writeStream.format("fake").start("/output_path")
 
+**Admission Control for Streaming Sources**
+
+Spark supports admission control for streaming sources to limit the amount of data processed in each micro-batch. This helps control resource usage and maintain consistent processing times. Python streaming data sources support three types of admission control options:
+
+- **maxRecordsPerBatch**: Limit the maximum number of rows per batch
+- **maxFilesPerBatch**: Limit the maximum number of files per batch
+- **maxBytesPerBatch**: Limit the maximum bytes per batch (in bytes)
+
+These options can be specified when reading from a streaming source:
+
+.. code-block:: python
+
+    # Limit to 1000 rows per batch
+    query = spark.readStream \
+        .format("fake") \
+        .option("maxRecordsPerBatch", "1000") \
+        .load() \
+        .writeStream \
+        .format("console") \
+        .start()
+
+    # Limit to 100 files per batch
+    query = spark.readStream \
+        .format("fake") \
+        .option("maxFilesPerBatch", "100") \
+        .load() \
+        .writeStream \
+        .format("console") \
+        .start()
+
+    # Limit to 10 MB per batch
+    query = spark.readStream \
+        .format("fake") \
+        .option("maxBytesPerBatch", str(10 * 1024 * 1024)) \
+        .load() \
+        .writeStream \
+        .format("console") \
+        .start()
+
+**Note**: 
+
+- Only one admission control option should be specified at a time
+- All admission control values must be positive integers. If an invalid value is provided (negative, zero, or non-numeric), an ``IllegalArgumentException`` will be thrown
+- These options apply to both ``streamReader()`` and ``simpleStreamReader()`` implementations
+- This behavior is consistent with Spark's built-in streaming sources (e.g., ``maxFilesPerTrigger``, ``maxBytesPerTrigger`` for file sources)
+
+**Implementing Admission Control in Custom Data Sources (Advanced)**
+
+For users implementing the full ``DataSourceStreamReader`` API (not ``SimpleDataSourceStreamReader``), admission control is enabled through optional parameters on the ``latestOffset()`` method. As of Spark 4.2, the ``latestOffset()`` method signature has been enhanced to accept optional parameters:
+
+.. code-block:: python
+
+    from pyspark.sql.datasource import DataSourceStreamReader
+    from typing import Optional, Union, Tuple
+    
+    class MyStreamReader(DataSourceStreamReader):
+        def latestOffset(
+            self, 
+            start: Optional[dict] = None, 
+            limit: Optional[dict] = None
+        ) -> Union[dict, Tuple[dict, dict]]:
+            # For backward compatibility: old implementations without parameters still work
+            if start is None and limit is None:
+                # Old behavior: return latest offset without admission control
+                return {"offset": self.get_latest_offset()}
+            
+            # New behavior: with admission control support
+            true_latest = self.get_latest_offset()
+            
+            # Apply admission control if configured
+            if limit and limit.get("type") == "maxRows":
+                max_rows = limit["maxRows"]
+                capped_offset = self.calculate_capped_offset(start, max_rows)
+                # Return tuple: (capped_offset, true_latest_offset)
+                return (capped_offset, true_latest)
+            elif limit and limit.get("type") == "maxFiles":
+                max_files = limit["maxFiles"]
+                capped_offset = self.calculate_capped_offset_by_files(start, max_files)
+                return (capped_offset, true_latest)
+            else:
+                # No limit or allAvailable
+                return (true_latest, true_latest)
+
+**Key Points:**
+
+- **Backward Compatibility**: Old implementations that don't accept parameters continue to work without modification
+- **Optional Parameters**: Both ``start`` and ``limit`` are optional; if not provided, implement old behavior
+- **Return Type**: Return a single ``dict`` for old behavior, or a ``Tuple[dict, dict]`` for new behavior with admission control
+- **Limit Structure**: The ``limit`` parameter is a dictionary with:
+  - ``{"type": "maxRows", "maxRows": N}`` for row-based limits
+  - ``{"type": "maxFiles", "maxFiles": N}`` for file-based limits  
+  - ``{"type": "maxBytes", "maxBytes": N}`` for byte-based limits
+  - ``{"type": "allAvailable"}`` for no limit
+- **SimpleDataSourceStreamReader**: Users of the simple API don't need to implement this; the framework handles admission control automatically
+
 Python Data Source Reader with direct Arrow Batch support for improved performance
 ----------------------------------------------------------------------------------
 The Python Datasource Reader supports direct yielding of Arrow Batches, which can significantly improve data processing performance. By using the efficient Arrow format,

python/pyspark/sql/datasource.py

@@ -714,20 +714,46 @@ def initialOffset(self) -> dict:
             messageParameters={"feature": "initialOffset"},
         )
 
-    def latestOffset(self) -> dict:
+    def latestOffset(
+        self, start: Optional[dict] = None, limit: Optional[dict] = None
+    ) -> Union[dict, Tuple[dict, dict]]:
         """
         Returns the most recent offset available.
 
+        Parameters (optional - added in Spark 4.2 for admission control)
+        -------------------------------------------------------------------
+        start : dict, optional
+            The starting offset. Enables admission control when provided.
+        limit : dict, optional
+            Admission control limit with structure:
+            - {"type": "maxRows", "maxRows": N}
+            - {"type": "maxFiles", "maxFiles": N}
+            - {"type": "maxBytes", "maxBytes": N}
+            - {"type": "allAvailable"}
+
         Returns
         -------
-        dict
-            A dict or recursive dict whose key and value are primitive types, which includes
-            Integer, String and Boolean.
+        dict or Tuple[dict, dict]
+            - Old behavior (no params): returns single offset dict
+            - New behavior (with params): returns (capped_offset, true_latest_offset)
 
         Examples
         --------
+        Old implementation (backward compatible):
+
         >>> def latestOffset(self):
-        ...     return {"parititon-1": {"index": 3, "closed": True}, "partition-2": {"index": 5}}
+        ...     return {"partition-1": {"index": 3, "closed": True}, "partition-2": {"index": 5}}
+
+        New implementation (with admission control):
+
+        >>> def latestOffset(self, start=None, limit=None):
+        ...     if start is None or limit is None:
+        ...         return {"offset": self.get_latest()}
+        ...     if limit["type"] == "maxRows":
+        ...         capped = self.calc_capped(start, limit["maxRows"])
+        ...         true_latest = self.get_latest()
+        ...         return (capped, true_latest)
+        ...     return (self.get_latest(), self.get_latest())
         """
         raise PySparkNotImplementedError(
             errorClass="NOT_IMPLEMENTED",

python/pyspark/sql/datasource_internal.py

@@ -19,7 +19,7 @@
 import json
 import copy
 from itertools import chain
-from typing import Iterator, List, Optional, Sequence, Tuple
+from typing import Iterator, List, Optional, Sequence, Tuple, Union
 
 from pyspark.sql.datasource import (
     DataSource,
@@ -88,14 +88,61 @@ def initialOffset(self) -> dict:
             self.initial_offset = self.simple_reader.initialOffset()
         return self.initial_offset
 
-    def latestOffset(self) -> dict:
+    def latestOffset(
+        self, start: Optional[dict] = None, limit: Optional[dict] = None
+    ) -> Union[dict, Tuple[dict, dict]]:
         # when query start for the first time, use initial offset as the start offset.
         if self.current_offset is None:
             self.current_offset = self.initialOffset()
-        (iter, end) = self.simple_reader.read(self.current_offset)
-        self.cache.append(PrefetchedCacheEntry(self.current_offset, end, iter))
-        self.current_offset = end
-        return end
+
+        # For backward compatibility: if called without parameters, use old behavior
+        if start is None and limit is None:
+            # Old behavior - no admission control
+            (full_iter, true_end) = self.simple_reader.read(self.current_offset)
+            self.cache.append(PrefetchedCacheEntry(self.current_offset, true_end, full_iter))
+            self.current_offset = true_end
+            return true_end
+
+        # New behavior with admission control support
+        # Call simple reader's read() to get all available data
+        (full_iter, true_end) = self.simple_reader.read(start)
+
+        # Check if admission control is enabled
+        if limit is not None and limit.get("type") == "maxRows":
+            max_rows = limit["maxRows"]
+            # Convert iterator to list to allow length calculation and slicing
+            data_list = list(full_iter)
+
+            if len(data_list) <= max_rows:
+                # All data fits within limit
+                capped_iter = iter(data_list)
+                capped_end = true_end
+            else:
+                # Cap the data to max_rows
+                capped_data = data_list[:max_rows]
+                capped_iter = iter(capped_data)
+                # Calculate capped offset based on how many rows we're actually taking
+                # For simple offset structures like {"offset": N}, we can calculate precisely
+                if (
+                    isinstance(start, dict)
+                    and "offset" in start
+                    and isinstance(true_end, dict)
+                    and "offset" in true_end
+                ):
+                    capped_end = {"offset": start["offset"] + len(capped_data)}
+                else:
+                    # For complex offsets, we can't reliably calculate partial offsets
+                    # User should use Full API (DataSourceStreamReader) for complex offset logic
+                    capped_end = true_end
+
+            self.cache.append(PrefetchedCacheEntry(start, capped_end, capped_iter))
+            self.current_offset = capped_end
+            return (capped_end, true_end)
+        else:
+            # No limit or allAvailable - return all data
+            self.cache.append(PrefetchedCacheEntry(start, true_end, full_iter))
+            self.current_offset = true_end
+            return (true_end, true_end)
 
     def commit(self, end: dict) -> None:
         if self.current_offset is None: