[WIP][SQL] Prototype a Rust SDK for native batch data sources

examples/src/main/rust/native-data-source/.gitignore

@@ -0,0 +1,2 @@
+*.sparkpkg
+__pycache__/

examples/src/main/rust/native-data-source/Cargo.toml

@@ -0,0 +1,37 @@
+# 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.
+
+[package]
+name = "spark-native-range-demo"
+workspace = "../../../../../sql/native/rust"
+version.workspace = true
+edition.workspace = true
+license.workspace = true
+publish.workspace = true
+
+[lib]
+name = "spark_datasource_rust_sdk_range"
+path = "src/lib.rs"
+crate-type = ["cdylib", "rlib"]
+
+[dependencies]
+spark-native-source = { path = "../../../../../sql/native/rust" }
+arrow-array.workspace = true
+arrow-schema.workspace = true
+serde.workspace = true
+serde_json.workspace = true
+
+[dev-dependencies]
+jni = { workspace = true, features = ["invocation"] }

examples/src/main/rust/native-data-source/README.md

@@ -0,0 +1,90 @@
+<!--
+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.
+-->
+
+# Rust native range data source
+
+This example implements batch reads using the [Rust native data source SDK][sdk].
+The connector in [src/lib.rs](src/lib.rs) exports `rust_sdk_range` through the SDK's JNI macro.
+See the SDK guide for API contracts, ownership rules, and local test instructions.
+
+The example targets NativeBridge ABI version 1 from [Spark PR #59313][native-source-pr], pinned
+to commit `ff7b1ae0150fbcb3b6a7c67bea8accf6aa4ad9bd`. Set `SPARK_HOME` to a Spark distribution
+built with that PR before running the Spark commands below.
+
+## Data and options
+
+| Column | Spark type | Value |
+| --- | --- | --- |
+| `id` | `BIGINT`, non-null | Integers in `[start, end)` |
+| `label` | `STRING`, non-null | `row-{id}`, such as `row-13` |
+
+| Option | Default | Meaning |
+| --- | --- | --- |
+| `start` | `0` | Inclusive start, as a signed 64-bit integer |
+| `end` | `100` | Exclusive end; must be at least `start` |
+| `numpartitions` | `2` | Maximum number of partitions; must be positive |
+| `batchsize` | `1024` | Maximum rows per Arrow batch; must be positive |
+
+Spark normalizes option names to lower case. Unknown options are rejected.
+
+The source evaluates `=`, `<=>`, `<`, `<=`, `>`, and `>=` comparisons between `id` and non-null
+integer literals, including conjunctions of supported predicates. Other predicates remain in
+Spark. A global limit is applied after exact filters and before splitting the range into partitions;
+the SDK declines limit pushdown when a residual predicate remains. Column pruning preserves the
+requested order and supports zero-column batches for operations such as `COUNT(*)`.
+
+## Build and run
+
+Run these commands from the repository root. The workspace, lockfile, and build output belong to
+`sql/native/rust`; the example is a member of that workspace.
+
+```sh
+cargo build --manifest-path sql/native/rust/Cargo.toml --workspace --locked
+
+python3 examples/src/main/rust/native-data-source/package.py \
+  --library sql/native/rust/target/debug/libspark_datasource_rust_sdk_range.so \
+  --output sql/native/rust/target/rust_sdk_range.sparkpkg
+
+"$SPARK_HOME/bin/spark-submit" --master 'local[2]' \
+  examples/src/main/rust/native-data-source/smoke.py \
+  --package sql/native/rust/target/rust_sdk_range.sparkpkg
+```
+
+The library name above is for Linux. On macOS use `libspark_datasource_rust_sdk_range.dylib`;
+on Windows use `spark_datasource_rust_sdk_range.dll`. The packaging script detects the host
+platform. For a cross-compiled library, supply `--platform`, for example `linux-aarch64`.
+
+The package can also be used from a Python session backed by the same Spark build:
+
+```python
+spark.addArtifact("sql/native/rust/target/rust_sdk_range.sparkpkg")
+df = (
+    spark.read.format("rust_sdk_range")
+    .option("end", 100)
+    .option("numpartitions", 4)
+    .load()
+)
+df.where("id >= 10 AND id < 20").select("label", "id").show()
+```
+
+[smoke.py](smoke.py) checks the resulting rows, projection, limits, zero-column counts, and
+filters that Spark must evaluate. This Spark smoke script has not yet been run against the
+required distribution. The local JVM integration test exercises JNI and Arrow C Stream without
+launching Spark; it does not establish end-to-end Spark compatibility.
+
+[sdk]: ../../../../../sql/native/rust/README.md
+[native-source-pr]: https://github.com/apache/spark/pull/59313

examples/src/main/rust/native-data-source/package.py

@@ -0,0 +1,70 @@
+#!/usr/bin/env python3
+# 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.
+
+"""Package the compiled range connector for the NativeBridge ABI in Spark PR #59313."""
+
+import argparse
+import json
+from pathlib import Path
+import platform
+import zipfile
+
+
+def current_platform():
+    os_name = {"Linux": "linux", "Darwin": "osx", "Windows": "windows"}.get(platform.system())
+    arch = {
+        "x86_64": "x86_64",
+        "amd64": "x86_64",
+        "aarch64": "aarch64",
+        "arm64": "aarch64",
+    }.get(platform.machine().lower())
+    return f"{os_name}-{arch}" if os_name and arch else None
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--library", type=Path, required=True)
+    parser.add_argument("--output", type=Path, required=True)
+    parser.add_argument("--platform", default=current_platform(), help="Override for cross builds")
+    args = parser.parse_args()
+    if not args.library.is_file():
+        parser.error(f"Library does not exist: {args.library}")
+    if args.output.suffix != ".sparkpkg":
+        parser.error("The output file must have the .sparkpkg extension")
+    if args.platform not in {
+        f"{os_name}-{arch}"
+        for os_name in ("linux", "osx", "windows")
+        for arch in ("x86_64", "aarch64")
+    }:
+        parser.error("Specify a supported --platform, such as linux-x86_64")
+    if args.library.resolve() == args.output.resolve():
+        parser.error("The output must differ from the input library")
+
+    library_path = f"{args.platform}/{args.library.name}"
+    manifest = {
+        "abiVersion": 1,
+        "dataSources": ["rust_sdk_range"],
+        "libraries": {args.platform: library_path},
+    }
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    with zipfile.ZipFile(args.output, "w", compression=zipfile.ZIP_DEFLATED) as package:
+        package.writestr("spark-native-datasource.json", json.dumps(manifest, indent=2) + "\n")
+        package.write(args.library, library_path)
+    print(args.output.resolve())
+
+
+if __name__ == "__main__":
+    main()

examples/src/main/rust/native-data-source/smoke.py

@@ -0,0 +1,71 @@
+#!/usr/bin/env python3
+# 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.
+
+"""Run with spark-submit from a Spark distribution that includes PR #59313."""
+
+import argparse
+from pathlib import Path
+
+from pyspark.sql import SparkSession, functions as F
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--package", type=Path, required=True)
+    args = parser.parse_args()
+    package = args.package.resolve(strict=True)
+
+    spark = SparkSession.builder.appName("Rust native data source smoke").getOrCreate()
+    try:
+        spark.addArtifact(str(package))
+        df = (
+            spark.read.format("rust_sdk_range")
+            .option("start", 0)
+            .option("end", 31)
+            .option("numpartitions", 4)
+            .option("batchsize", 3)
+            .load()
+        )
+        assert [tuple(row) for row in df.orderBy("id").collect()] == [
+            (i, f"row-{i}") for i in range(31)
+        ]
+
+        projected = df.where((F.col("id") >= 7) & (F.col("id") < 17)).select("label", "id")
+        assert [tuple(row) for row in projected.orderBy("id").collect()] == [
+            (f"row-{i}", i) for i in range(7, 17)
+        ]
+        limited = df.where((F.col("id") >= 10) & (F.col("id") < 20)).limit(4).collect()
+        assert len(limited) == 4
+        assert all(10 <= row.id < 20 for row in limited)
+
+        # Count needs row counts even when column pruning produces zero-column Arrow batches.
+        assert df.select().count() == 31
+        assert df.where(F.col("id") < 0).count() == 0
+
+        # The connector does not evaluate these expressions; Spark must retain the filters.
+        arithmetic = df.where((F.col("id") % 3) == 1).select("id").orderBy("id").collect()
+        assert [row.id for row in arithmetic] == list(range(1, 31, 3))
+        labels = df.where(F.col("label").endswith("3")).select("id").orderBy("id").collect()
+        assert [row.id for row in labels] == [3, 13, 23]
+        residual_limit = df.where(F.col("label") == "row-21").limit(1).select("id").collect()
+        assert [row.id for row in residual_limit] == [21]
+        print("Rust native data source Spark smoke checks passed")
+    finally:
+        spark.stop()
+
+
+if __name__ == "__main__":
+    main()