From 6c4a4d3ec95838602a1dc8bdd23eb222fd13bc17 Mon Sep 17 00:00:00 2001 From: xuanyili Date: Sat, 5 Sep 2026 23:09:43 +0000 Subject: [PATCH 1/6] feat: add Lance contrib build gate --- contrib/lance/native/Cargo.toml | 32 +++++++++ contrib/lance/native/src/lib.rs | 18 +++++ contrib/lance/native/src/planner.rs | 33 ++++++++++ native/Cargo.lock | 9 +++ native/Cargo.toml | 5 +- native/core/Cargo.toml | 3 + native/core/src/execution/planner.rs | 10 ++- .../core/src/execution/planner/lance_scan.rs | 66 +++++++++++++++++++ native/proto/src/proto/operator.proto | 26 ++++++++ 9 files changed, 197 insertions(+), 5 deletions(-) create mode 100644 contrib/lance/native/Cargo.toml create mode 100644 contrib/lance/native/src/lib.rs create mode 100644 contrib/lance/native/src/planner.rs create mode 100644 native/core/src/execution/planner/lance_scan.rs diff --git a/contrib/lance/native/Cargo.toml b/contrib/lance/native/Cargo.toml new file mode 100644 index 00000000000..e34400c6a67 --- /dev/null +++ b/contrib/lance/native/Cargo.toml @@ -0,0 +1,32 @@ +# 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 = "comet-contrib-lance" +description = "Native Lance scan integration for Apache DataFusion Comet" +version = "1.1.0" +edition = "2021" +rust-version = "1.94" +publish = false +license = "Apache-2.0" + +[lib] +crate-type = ["rlib"] + +[dependencies] +datafusion = { version = "55.0.0", default-features = false } +datafusion-comet-proto = { path = "../../../native/proto" } diff --git a/contrib/lance/native/src/lib.rs b/contrib/lance/native/src/lib.rs new file mode 100644 index 00000000000..219efeac515 --- /dev/null +++ b/contrib/lance/native/src/lib.rs @@ -0,0 +1,18 @@ +// 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. + +pub mod planner; diff --git a/contrib/lance/native/src/planner.rs b/contrib/lance/native/src/planner.rs new file mode 100644 index 00000000000..1d5a460dd02 --- /dev/null +++ b/contrib/lance/native/src/planner.rs @@ -0,0 +1,33 @@ +// 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. + +use std::sync::Arc; + +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::error::DataFusionError; +use datafusion::physical_plan::ExecutionPlan; +use datafusion_comet_proto::spark_operator::{LanceScanCommon, LanceScanPartition}; + +pub fn plan_lance_scan( + _common: &LanceScanCommon, + _partition: &LanceScanPartition, + _output_schema: &SchemaRef, +) -> Result, DataFusionError> { + Err(DataFusionError::NotImplemented( + "contrib-lance native Lance read path is not yet implemented (build-gate stub)".to_string(), + )) +} diff --git a/native/Cargo.lock b/native/Cargo.lock index df5fd4ac14a..fabb56e2ec8 100644 --- a/native/Cargo.lock +++ b/native/Cargo.lock @@ -1480,6 +1480,14 @@ dependencies = [ "datafusion-comet-proto", ] +[[package]] +name = "comet-contrib-lance" +version = "1.1.0" +dependencies = [ + "datafusion", + "datafusion-comet-proto", +] + [[package]] name = "comfy-table" version = "7.2.2" @@ -1966,6 +1974,7 @@ dependencies = [ "base64 0.23.1", "bytes", "comet-contrib-delta", + "comet-contrib-lance", "criterion", "datafusion", "datafusion-comet-common", diff --git a/native/Cargo.toml b/native/Cargo.toml index 1805a185e9d..bde2ba432c0 100644 --- a/native/Cargo.toml +++ b/native/Cargo.toml @@ -18,9 +18,8 @@ [workspace] default-members = ["core", "spark-expr", "common", "proto", "jni-bridge", "shuffle"] members = ["core", "spark-expr", "common", "proto", "jni-bridge", "shuffle"] -# The contrib crate at ../contrib/delta/native is intentionally NOT a workspace member -# (workspace members must live hierarchically under the workspace root). It's pulled in -# as a path dep by `core/Cargo.toml` when the `contrib-delta` feature is enabled. +# Crates under ../contrib are intentionally NOT workspace members. Core pulls them in as +# optional path dependencies when their corresponding contrib feature is enabled. exclude = ["../contrib"] resolver = "2" diff --git a/native/core/Cargo.toml b/native/core/Cargo.toml index 97ce3808cea..f0c7735a503 100644 --- a/native/core/Cargo.toml +++ b/native/core/Cargo.toml @@ -72,6 +72,8 @@ parking_lot = "0.12.5" # Optional Delta Lake contrib (enabled by the `contrib-delta` feature). Source lives # under `contrib/delta/native/` so non-Delta committers can ignore it. comet-contrib-delta = { path = "../../contrib/delta/native", optional = true } +# Optional Lance contrib. The native scan implementation lives outside core. +comet-contrib-lance = { path = "../../contrib/lance/native", optional = true } reqwest = { version = "0.12", default-features = false, features = ["rustls-tls-native-roots", "http2"] } object_store_opendal = { version = "0.58.0", optional = true } hdfs-sys = {version = "0.3", optional = true, features = ["hdfs_3_3"]} @@ -104,6 +106,7 @@ datafusion-functions-nested = { version = "55.1.0" } [features] backtrace = ["datafusion/backtrace"] default = ["hdfs-opendal"] +contrib-lance = ["dep:comet-contrib-lance"] hdfs-opendal = ["opendal", "object_store_opendal", "hdfs-sys"] jemalloc = ["tikv-jemallocator", "tikv-jemalloc-ctl"] # Delta Lake integration. When enabled, links the `comet-contrib-delta` crate diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 5bd4bfebb77..60c206617e3 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -29,6 +29,8 @@ pub mod operator_registry; // and calls into that crate. #[cfg(feature = "contrib-delta")] mod delta_scan; +#[cfg(feature = "contrib-lance")] +mod lance_scan; use crate::execution::operators::init_csv_datasource_exec; use crate::execution::operators::AlignedArrowStreamReader; @@ -1876,10 +1878,14 @@ impl PhysicalPlanner { if let Some(result) = delta_scan::try_plan_contrib_scan(self, spark_plan, contrib) { return result; } + #[cfg(feature = "contrib-lance")] + if let Some(result) = lance_scan::try_plan_contrib_scan(self, spark_plan, contrib) { + return result; + } Err(GeneralError(format!( "Received a contrib_scan operator (type_url: {}) but core was built without a \ - contrib that handles it. Rebuild with the matching contrib feature -- e.g. \ - `-Pcontrib-delta` (Maven) + `--features contrib-delta` (Cargo) for Delta Lake.", + contrib that handles it. Rebuild with the matching Maven profile and Cargo \ + feature.", contrib.type_url ))) } diff --git a/native/core/src/execution/planner/lance_scan.rs b/native/core/src/execution/planner/lance_scan.rs new file mode 100644 index 00000000000..c0a46d25c57 --- /dev/null +++ b/native/core/src/execution/planner/lance_scan.rs @@ -0,0 +1,66 @@ +// 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. + +use std::sync::Arc; + +use datafusion_comet_proto::spark_operator::{ContribScan, LanceScan, Operator}; +use prost::Message; + +use crate::execution::operators::ExecutionError::GeneralError; +use crate::execution::planner::{ + convert_spark_types_to_arrow_schema, PhysicalPlanner, PlanCreationResult, +}; +use crate::execution::spark_plan::SparkPlan; + +const LANCE_SCAN_TYPE_NAME: &str = "comet.contrib.lance.LanceScan"; + +pub(crate) fn try_plan_contrib_scan( + _planner: &PhysicalPlanner, + spark_plan: &Operator, + contrib: &ContribScan, +) -> Option { + if !contrib.type_url.ends_with(LANCE_SCAN_TYPE_NAME) { + return None; + } + + Some( + LanceScan::decode(contrib.value.as_slice()) + .map_err(|e| GeneralError(format!("Failed to decode LanceScan: {e}"))) + .and_then(|scan| plan_lance_scan(spark_plan, &scan)), + ) +} + +fn plan_lance_scan(spark_plan: &Operator, scan: &LanceScan) -> PlanCreationResult { + let common = scan + .common + .as_ref() + .ok_or_else(|| GeneralError("LanceScan missing common data".into()))?; + let partition = scan + .partition + .as_ref() + .ok_or_else(|| GeneralError("LanceScan missing partition data".into()))?; + let output_schema = convert_spark_types_to_arrow_schema(common.projected_schema.as_slice()); + + let exec = comet_contrib_lance::planner::plan_lance_scan(common, partition, &output_schema) + .map_err(|e| GeneralError(e.to_string()))?; + + Ok(( + vec![], + vec![], + Arc::new(SparkPlan::new(spark_plan.plan_id, exec, vec![])), + )) +} diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 2a89ded3507..819806213df 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -356,6 +356,32 @@ message IcebergScan { repeated IcebergFileScanTask file_scan_tasks = 2; } +// Common data shared by all partitions for native Lance scans. +message LanceScanCommon { + string scan_id = 1; + repeated SparkStructField required_schema = 2; + string native_scan_plan_class = 3; + string dataset_uri = 4; + int64 resolved_version = 5; + map storage_options = 6; + repeated SparkStructField projected_schema = 7; + optional string filter_sql = 8; + optional int64 limit = 9; + optional int64 offset = 10; + uint32 batch_size = 11; + uint32 descriptor_version = 12; +} + +message LanceScanPartition { + uint32 partition_index = 1; + repeated uint32 fragment_ids = 2; +} + +message LanceScan { + LanceScanCommon common = 1; + LanceScanPartition partition = 2; +} + // Helper message for deduplicating field ID lists message ProjectFieldIdList { repeated int32 field_ids = 1; From ac7b6bb3d5b884a6014b97923f762220de163bf9 Mon Sep 17 00:00:00 2001 From: xuanyili Date: Thu, 11 Jun 2026 07:31:35 +0000 Subject: [PATCH 2/6] feat: scaffold contrib Lance native scan --- pom.xml | 7 + spark/pom.xml | 25 +++ .../comet/lance/CometLanceSupport.scala | 52 +++++ .../serde/operator/CometLanceNativeScan.scala | 119 ++++++++++++ .../sql/comet/CometLanceNativeScanExec.scala | 116 ++++++++++++ .../scala/org/apache/comet/CometConf.scala | 8 + .../apache/comet/lance/LanceIntegration.scala | 178 ++++++++++++++++++ .../sql/comet/CometLanceNativeScanLike.scala | 28 +++ .../comet/rules/CometScanRuleSuite.scala | 4 + 9 files changed, 537 insertions(+) create mode 100644 spark/src/contrib-lance/scala/org/apache/comet/lance/CometLanceSupport.scala create mode 100644 spark/src/contrib-lance/scala/org/apache/comet/serde/operator/CometLanceNativeScan.scala create mode 100644 spark/src/contrib-lance/scala/org/apache/spark/sql/comet/CometLanceNativeScanExec.scala create mode 100644 spark/src/main/scala/org/apache/comet/lance/LanceIntegration.scala create mode 100644 spark/src/main/scala/org/apache/spark/sql/comet/CometLanceNativeScanLike.scala diff --git a/pom.xml b/pom.xml index f4b2be220ec..9634dd0cffc 100644 --- a/pom.xml +++ b/pom.xml @@ -762,6 +762,13 @@ under the License. + + contrib-lance + + true + + + scala-2.12 diff --git a/spark/pom.xml b/spark/pom.xml index 8dc632d5b92..19341cd5127 100644 --- a/spark/pom.xml +++ b/spark/pom.xml @@ -505,6 +505,31 @@ under the License. + + contrib-lance + + + + org.codehaus.mojo + build-helper-maven-plugin + + + add-contrib-lance-source + generate-sources + + add-source + + + + src/contrib-lance/scala + + + + + + + + generate-docs diff --git a/spark/src/contrib-lance/scala/org/apache/comet/lance/CometLanceSupport.scala b/spark/src/contrib-lance/scala/org/apache/comet/lance/CometLanceSupport.scala new file mode 100644 index 00000000000..d02f73dec5d --- /dev/null +++ b/spark/src/contrib-lance/scala/org/apache/comet/lance/CometLanceSupport.scala @@ -0,0 +1,52 @@ +/* + * 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 org.apache.comet.lance + +import scala.collection.mutable.ListBuffer + +import org.apache.spark.sql.comet.CometBatchScanExec +import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.datasources.v2.BatchScanExec + +import org.apache.comet.CometSparkSessionExtensions.withInfos +import org.apache.comet.serde.OperatorOuterClass.Operator +import org.apache.comet.serde.operator.CometLanceNativeScan + +object CometLanceSupport { + + def tryTransform(scanExec: BatchScanExec, nativeScanPlan: Object): Option[SparkPlan] = { + val fallbackReasons = new ListBuffer[String]() + val schemaSupported = + CometBatchScanExec.isSchemaSupported(scanExec.scan.readSchema(), fallbackReasons) + + if (!schemaSupported) { + fallbackReasons += s"Schema ${scanExec.scan.readSchema()} is not supported" + withInfos(scanExec, fallbackReasons.toSet) + None + } else { + val builder = Operator.newBuilder().setPlanId(scanExec.id) + CometLanceNativeScan + .convert(scanExec, builder, Option(nativeScanPlan)) + .map { nativeOp => + CometLanceNativeScan.createExec(nativeOp, scanExec, Option(nativeScanPlan)) + } + } + } +} diff --git a/spark/src/contrib-lance/scala/org/apache/comet/serde/operator/CometLanceNativeScan.scala b/spark/src/contrib-lance/scala/org/apache/comet/serde/operator/CometLanceNativeScan.scala new file mode 100644 index 00000000000..b17a9f5df5b --- /dev/null +++ b/spark/src/contrib-lance/scala/org/apache/comet/serde/operator/CometLanceNativeScan.scala @@ -0,0 +1,119 @@ +/* + * 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 org.apache.comet.serde.operator + +import scala.jdk.CollectionConverters._ + +import org.apache.spark.internal.Logging +import org.apache.spark.sql.comet.{CometLanceNativeScanExec, CometNativeExec, SerializedPlan} +import org.apache.spark.sql.execution.datasources.v2.BatchScanExec +import org.apache.spark.sql.types.StructType + +import org.apache.comet.{CometConf, ConfigEntry} +import org.apache.comet.serde.{CometOperatorSerde, Compatible, OperatorOuterClass, SupportLevel} +import org.apache.comet.serde.OperatorOuterClass.Operator + +object CometLanceNativeScan extends CometOperatorSerde[BatchScanExec] with Logging { + + override def enabledConfig: Option[ConfigEntry[Boolean]] = + Some(CometConf.COMET_LANCE_NATIVE_ENABLED) + + override def getSupportLevel(operator: BatchScanExec): SupportLevel = Compatible() + + override def convert( + scanExec: BatchScanExec, + builder: Operator.Builder, + childOp: Operator*): Option[Operator] = + convert(scanExec, builder, None) + + def convert( + scanExec: BatchScanExec, + builder: Operator.Builder, + nativeScanPlan: Option[Any]): Option[Operator] = { + val sourceKey = scanKey(scanExec) + val requiredSchema = scanExec.scan.readSchema() + val nativePlanClass = nativeScanPlan + .map(_.getClass.getName) + .getOrElse("") + + val commonBuilder = OperatorOuterClass.LanceScanCommon + .newBuilder() + .setScanId(sourceKey) + .setNativeScanPlanClass(nativePlanClass) + .addAllRequiredSchema(schema2Proto(requiredSchema.fields).toSeq.asJava) + + val lanceScanBuilder = OperatorOuterClass.LanceScan + .newBuilder() + .setCommon(commonBuilder.build()) + + builder.clearChildren() + Some(builder.setLanceScan(lanceScanBuilder).build()) + } + + override def createExec(nativeOp: Operator, op: BatchScanExec): CometNativeExec = + createExec(nativeOp, op, None) + + def createExec( + nativeOp: Operator, + op: BatchScanExec, + nativeScanPlan: Option[Any]): CometNativeExec = { + val nativePlanClass = nativeScanPlan + .map(_.getClass.getName) + .getOrElse("") + val exec = CometLanceNativeScanExec( + nativeOp, + op.output, + op.runtimeFilters, + op.scan.readSchema(), + op, + SerializedPlan(None), + scanKey(op), + nativePlanClass) + op.logicalLink.foreach(exec.setLogicalLink) + exec + } + + def serializePartitions( + sourceKey: String, + requiredSchema: StructType, + nativeScanPlanClassName: String): (Array[Byte], Array[Array[Byte]]) = { + val common = OperatorOuterClass.LanceScanCommon + .newBuilder() + .setScanId(sourceKey) + .setNativeScanPlanClass(nativeScanPlanClassName) + .addAllRequiredSchema(schema2Proto(requiredSchema.fields).toSeq.asJava) + .build() + + val partition = OperatorOuterClass.LanceScanPartition + .newBuilder() + .setPartitionIndex(0) + .build() + + val partitionScan = OperatorOuterClass.LanceScan + .newBuilder() + .setPartition(partition) + .build() + + (common.toByteArray, Array(partitionScan.toByteArray)) + } + + private def scanKey(scanExec: BatchScanExec): String = + s"lance_${scanExec.id}_${scanExec.scan.hashCode()}" +} diff --git a/spark/src/contrib-lance/scala/org/apache/spark/sql/comet/CometLanceNativeScanExec.scala b/spark/src/contrib-lance/scala/org/apache/spark/sql/comet/CometLanceNativeScanExec.scala new file mode 100644 index 00000000000..da82e06e421 --- /dev/null +++ b/spark/src/contrib-lance/scala/org/apache/spark/sql/comet/CometLanceNativeScanExec.scala @@ -0,0 +1,116 @@ +/* + * 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 org.apache.spark.sql.comet + +import scala.jdk.CollectionConverters._ + +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, SortOrder} +import org.apache.spark.sql.catalyst.plans.QueryPlan +import org.apache.spark.sql.catalyst.plans.physical.{Partitioning, UnknownPartitioning} +import org.apache.spark.sql.execution.datasources.v2.BatchScanExec +import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.vectorized.ColumnarBatch + +import com.google.common.base.Objects + +import org.apache.comet.serde.OperatorOuterClass.Operator +import org.apache.comet.serde.operator.CometLanceNativeScan + +case class CometLanceNativeScanExec( + override val nativeOp: Operator, + override val output: Seq[Attribute], + runtimeFilters: Seq[Expression], + requiredSchema: StructType, + @transient originalPlan: BatchScanExec, + override val serializedPlanOpt: SerializedPlan, + override val sourceKey: String, + nativeScanPlanClassName: String) + extends CometLeafExec + with CometLanceNativeScanLike { + + override val supportsColumnar: Boolean = true + + override val nodeName: String = "CometLanceNativeScan" + + @transient private lazy val serializedPartitionData: (Array[Byte], Array[Array[Byte]]) = + CometLanceNativeScan.serializePartitions(sourceKey, requiredSchema, nativeScanPlanClassName) + + override def commonData: Array[Byte] = serializedPartitionData._1 + + override def perPartitionData: Array[Array[Byte]] = serializedPartitionData._2 + + override lazy val outputPartitioning: Partitioning = + UnknownPartitioning(perPartitionData.length) + + override lazy val outputOrdering: Seq[SortOrder] = Nil + + override def doExecuteColumnar(): RDD[ColumnarBatch] = { + throw new UnsupportedOperationException( + "Native Lance scan execution is not implemented yet") + } + + override def convertBlock(): CometLanceNativeScanExec = { + val newSerializedPlan = if (serializedPlanOpt.isEmpty) { + SerializedPlan(Some(CometExec.serializeNativePlan(nativeOp))) + } else { + serializedPlanOpt + } + + CometLanceNativeScanExec( + nativeOp, + output, + runtimeFilters, + requiredSchema, + originalPlan, + newSerializedPlan, + sourceKey, + nativeScanPlanClassName) + } + + override protected def doCanonicalize(): CometLanceNativeScanExec = { + CometLanceNativeScanExec( + nativeOp, + output.map(QueryPlan.normalizeExpressions(_, output)), + QueryPlan.normalizePredicates( + CometScanUtils.filterUnusedDynamicPruningExpressions(runtimeFilters), + output), + requiredSchema, + null, + SerializedPlan(None), + sourceKey, + nativeScanPlanClassName) + } + + override def stringArgs: Iterator[Any] = + Iterator(output, s"$sourceKey, nativeScanPlan=$nativeScanPlanClassName") + + override def equals(obj: Any): Boolean = obj match { + case other: CometLanceNativeScanExec => + this.sourceKey == other.sourceKey && + this.output == other.output && + this.runtimeFilters == other.runtimeFilters && + this.serializedPlanOpt == other.serializedPlanOpt + case _ => false + } + + override def hashCode(): Int = + Objects.hashCode(sourceKey, output.asJava, runtimeFilters, serializedPlanOpt) +} diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 43f030a7d9d..6068527d51a 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -141,6 +141,14 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(false) + val COMET_LANCE_NATIVE_ENABLED: ConfigEntry[Boolean] = + conf("spark.comet.scan.lanceNative.enabled") + .category(CATEGORY_SCAN) + .doc("Whether to enable native Lance table scans through the optional contrib-lance " + + "integration. This is an experimental scaffold and is disabled by default.") + .booleanConf + .createWithDefault(false) + val COMET_ICEBERG_DATA_FILE_CONCURRENCY_LIMIT: ConfigEntry[Int] = conf("spark.comet.scan.icebergNative.dataFileConcurrencyLimit") .category(CATEGORY_SCAN) diff --git a/spark/src/main/scala/org/apache/comet/lance/LanceIntegration.scala b/spark/src/main/scala/org/apache/comet/lance/LanceIntegration.scala new file mode 100644 index 00000000000..348cb799f9c --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/lance/LanceIntegration.scala @@ -0,0 +1,178 @@ +/* + * 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 org.apache.comet.lance + +import java.lang.reflect.InvocationTargetException + +import scala.util.control.NonFatal + +import org.apache.spark.internal.Logging +import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.datasources.v2.BatchScanExec + +import org.apache.comet.CometConf +import org.apache.comet.CometSparkSessionExtensions.withInfo + +/** + * Reflection-only bridge for optional Lance Spark integration. + * + * Default Comet builds must not depend on Lance classes. This object treats both Lance Spark and + * the Comet contrib-lance scaffold as optional runtime classes and falls back cleanly when either + * side is absent. + */ +object LanceIntegration extends Logging { + + private val LanceScanClassName = "org.lance.spark.read.LanceScan" + private val NativeScanPlanMethod = "nativeScanPlan" + private val ContribSupportModule = "org.apache.comet.lance.CometLanceSupport$" + + def isLanceScan(scan: Any): Boolean = { + scan != null && { + scan.getClass.getName == LanceScanClassName || + loadClass(LanceScanClassName).exists(_.isInstance(scan)) + } + } + + def nativeScanPlan(scan: Any): Option[Any] = + if (isLanceScan(scan)) { + invokeNativeScanPlan(scan) + } else { + None + } + + def tryCreateNativeScan(scanExec: BatchScanExec): Option[SparkPlan] = { + if (!CometConf.COMET_LANCE_NATIVE_ENABLED.get(scanExec.conf)) { + withInfo( + scanExec, + s"Native Lance scan disabled because ${CometConf.COMET_LANCE_NATIVE_ENABLED.key} " + + "is not enabled") + return None + } + + if (!CometConf.COMET_EXEC_ENABLED.get(scanExec.conf)) { + withInfo( + scanExec, + s"Native Lance scan disabled because ${CometConf.COMET_EXEC_ENABLED.key} is not enabled") + return None + } + + val nativePlan = nativeScanPlan(scanExec.scan) match { + case Some(plan) => plan + case None => + withInfo( + scanExec, + s"Native Lance scan disabled because $LanceScanClassName.$NativeScanPlanMethod() " + + "is not available") + return None + } + + val support = loadContribSupport match { + case Some(module) => module + case None => + withInfo( + scanExec, + "Native Lance scan disabled because the contrib-lance build profile is not present") + return None + } + + try { + val method = + support.getClass.getMethod("tryTransform", classOf[BatchScanExec], classOf[Object]) + method.invoke(support, scanExec, nativePlan.asInstanceOf[AnyRef]) match { + case plan: Option[_] => plan.asInstanceOf[Option[SparkPlan]] + case other => + logWarning( + "Native Lance scan disabled because contrib-lance returned unexpected " + + s"result: ${Option(other).map(_.getClass.getName).getOrElse("null")}") + None + } + } catch { + case NonFatal(e) => + logWarning(s"Native Lance scan disabled by contrib-lance reflection failure: $e") + None + } + } + + private[comet] def invokeNativeScanPlan(scan: Any): Option[Any] = { + try { + findNoArgMethod(scan.getClass, NativeScanPlanMethod) + .flatMap { method => + Option(method.invoke(scan)) + } + } catch { + case e: InvocationTargetException => + logWarning( + s"Native Lance scan disabled because $NativeScanPlanMethod() threw: " + + s"${Option(e.getCause).map(_.getMessage).getOrElse(e.getMessage)}") + None + case NonFatal(e) => + logWarning(s"Native Lance scan disabled by reflection failure: $e") + None + } + } + + private def findNoArgMethod( + clazz: Class[_], + methodName: String): Option[java.lang.reflect.Method] = { + var current = clazz + while (current != null) { + try { + val method = current.getDeclaredMethod(methodName) + method.setAccessible(true) + return Some(method) + } catch { + case _: NoSuchMethodException => + current = current.getSuperclass + case NonFatal(_) => + return None + } + } + None + } + + private def loadContribSupport: Option[AnyRef] = + loadClass(ContribSupportModule).flatMap { clazz => + try { + Some(clazz.getField("MODULE$").get(null).asInstanceOf[AnyRef]) + } catch { + case NonFatal(_) => None + } + } + + private def loadClass(className: String): Option[Class[_]] = { + try { + val classLoader = Thread.currentThread().getContextClassLoader + // scalastyle:off classforname + val clazz = + if (classLoader != null) { + Class.forName(className, false, classLoader) + } else { + Class.forName(className) + } + // scalastyle:on classforname + Some(clazz) + } catch { + case _: ClassNotFoundException | _: NoClassDefFoundError => None + case NonFatal(e) => + logDebug(s"Unable to load optional class $className", e) + None + } + } +} diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometLanceNativeScanLike.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometLanceNativeScanLike.scala new file mode 100644 index 00000000000..72b5326b7b9 --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometLanceNativeScanLike.scala @@ -0,0 +1,28 @@ +/* + * 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 org.apache.spark.sql.comet + +private[comet] trait CometLanceNativeScanLike extends CometLeafExec { + def sourceKey: String + + def commonData: Array[Byte] + + def perPartitionData: Array[Array[Byte]] +} diff --git a/spark/src/test/scala/org/apache/comet/rules/CometScanRuleSuite.scala b/spark/src/test/scala/org/apache/comet/rules/CometScanRuleSuite.scala index b0211edf5cf..521afa44a9d 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometScanRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometScanRuleSuite.scala @@ -181,4 +181,8 @@ class CometScanRuleSuite extends CometTestBase { } } + test("Lance native scan config defaults to disabled") { + assert(!CometConf.COMET_LANCE_NATIVE_ENABLED.get()) + } + } From 893db78ba121a77ed6baf1b8e48d63e2941f32c6 Mon Sep 17 00:00:00 2001 From: xuanyili Date: Thu, 11 Jun 2026 08:08:59 +0000 Subject: [PATCH 3/6] feat: serialize Lance native scan descriptor --- native/proto/src/proto/operator.proto | 18 ++ .../serde/operator/CometLanceNativeScan.scala | 296 +++++++++++++++--- .../sql/comet/CometLanceNativeScanExec.scala | 36 ++- .../comet/rules/CometScanRuleSuite.scala | 108 +++++++ 4 files changed, 410 insertions(+), 48 deletions(-) diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 819806213df..2c2d55d385e 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -358,20 +358,38 @@ message IcebergScan { // Common data shared by all partitions for native Lance scans. message LanceScanCommon { + // Stable key for matching this scan during per-partition data injection. string scan_id = 1; + + // Schema Spark requested from the Lance scan. repeated SparkStructField required_schema = 2; + + // Informational class name of Lance's native scan plan object. string native_scan_plan_class = 3; + + // Lance dataset location and version resolved during Spark-side planning. string dataset_uri = 4; int64 resolved_version = 5; + + // Storage options captured by Lance Spark planning. map storage_options = 6; + + // Schema after projection pruning, as planned by Lance Spark. repeated SparkStructField projected_schema = 7; + + // Optional pushdowns captured by Lance Spark planning. optional string filter_sql = 8; optional int64 limit = 9; optional int64 offset = 10; + + // Preferred native batch size. uint32 batch_size = 11; + + // Version of the Lance Spark native scan descriptor contract. uint32 descriptor_version = 12; } +// Per-partition Lance scan split. message LanceScanPartition { uint32 partition_index = 1; repeated uint32 fragment_ids = 2; diff --git a/spark/src/contrib-lance/scala/org/apache/comet/serde/operator/CometLanceNativeScan.scala b/spark/src/contrib-lance/scala/org/apache/comet/serde/operator/CometLanceNativeScan.scala index b17a9f5df5b..e20ecf1df89 100644 --- a/spark/src/contrib-lance/scala/org/apache/comet/serde/operator/CometLanceNativeScan.scala +++ b/spark/src/contrib-lance/scala/org/apache/comet/serde/operator/CometLanceNativeScan.scala @@ -19,12 +19,15 @@ package org.apache.comet.serde.operator +import java.lang.reflect.InvocationTargetException + import scala.jdk.CollectionConverters._ +import scala.util.control.NonFatal import org.apache.spark.internal.Logging import org.apache.spark.sql.comet.{CometLanceNativeScanExec, CometNativeExec, SerializedPlan} import org.apache.spark.sql.execution.datasources.v2.BatchScanExec -import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.types.{DataType, StructType} import org.apache.comet.{CometConf, ConfigEntry} import org.apache.comet.serde.{CometOperatorSerde, Compatible, OperatorOuterClass, SupportLevel} @@ -32,6 +35,23 @@ import org.apache.comet.serde.OperatorOuterClass.Operator object CometLanceNativeScan extends CometOperatorSerde[BatchScanExec] with Logging { + case class LanceNativeScanSplitDescriptor(partitionIndex: Int, fragmentIds: Seq[Int]) + + case class LanceNativeScanDescriptor( + descriptorVersion: Int, + scanId: String, + datasetUri: String, + resolvedVersion: Long, + storageOptions: Map[String, String], + requiredSchema: StructType, + projectedSchema: StructType, + filterSql: Option[String], + limit: Option[Long], + offset: Option[Long], + batchSize: Int, + nativeScanPlanClass: String, + splits: Seq[LanceNativeScanSplitDescriptor]) + override def enabledConfig: Option[ConfigEntry[Boolean]] = Some(CometConf.COMET_LANCE_NATIVE_ENABLED) @@ -47,21 +67,11 @@ object CometLanceNativeScan extends CometOperatorSerde[BatchScanExec] with Loggi scanExec: BatchScanExec, builder: Operator.Builder, nativeScanPlan: Option[Any]): Option[Operator] = { - val sourceKey = scanKey(scanExec) - val requiredSchema = scanExec.scan.readSchema() - val nativePlanClass = nativeScanPlan - .map(_.getClass.getName) - .getOrElse("") - - val commonBuilder = OperatorOuterClass.LanceScanCommon - .newBuilder() - .setScanId(sourceKey) - .setNativeScanPlanClass(nativePlanClass) - .addAllRequiredSchema(schema2Proto(requiredSchema.fields).toSeq.asJava) + val descriptor = descriptorFor(scanExec, nativeScanPlan) val lanceScanBuilder = OperatorOuterClass.LanceScan .newBuilder() - .setCommon(commonBuilder.build()) + .setCommon(commonFromDescriptor(descriptor)) builder.clearChildren() Some(builder.setLanceScan(lanceScanBuilder).build()) @@ -74,46 +84,258 @@ object CometLanceNativeScan extends CometOperatorSerde[BatchScanExec] with Loggi nativeOp: Operator, op: BatchScanExec, nativeScanPlan: Option[Any]): CometNativeExec = { - val nativePlanClass = nativeScanPlan - .map(_.getClass.getName) - .getOrElse("") + val descriptor = descriptorFor(op, nativeScanPlan) val exec = CometLanceNativeScanExec( nativeOp, op.output, op.runtimeFilters, - op.scan.readSchema(), op, SerializedPlan(None), - scanKey(op), - nativePlanClass) + descriptor.scanId, + descriptor) op.logicalLink.foreach(exec.setLogicalLink) exec } - def serializePartitions( - sourceKey: String, - requiredSchema: StructType, - nativeScanPlanClassName: String): (Array[Byte], Array[Array[Byte]]) = { - val common = OperatorOuterClass.LanceScanCommon - .newBuilder() - .setScanId(sourceKey) - .setNativeScanPlanClass(nativeScanPlanClassName) - .addAllRequiredSchema(schema2Proto(requiredSchema.fields).toSeq.asJava) - .build() + def serializePartitions(descriptor: LanceNativeScanDescriptor): (Array[Byte], Array[Array[Byte]]) = + ( + commonFromDescriptor(descriptor).toByteArray, + descriptor.splits.map { split => + val partition = OperatorOuterClass.LanceScanPartition + .newBuilder() + .setPartitionIndex(split.partitionIndex) + .addAllFragmentIds(split.fragmentIds.map(Int.box).asJava) + .build() - val partition = OperatorOuterClass.LanceScanPartition - .newBuilder() - .setPartitionIndex(0) - .build() + OperatorOuterClass.LanceScan + .newBuilder() + .setPartition(partition) + .build() + .toByteArray + }.toArray) + + private[comet] def serializeNativePlan( + nativeScanPlan: Any, + fallbackScanId: String, + fallbackRequiredSchema: StructType): (Array[Byte], Array[Array[Byte]]) = { + serializePartitions(descriptorFromNativePlan( + nativeScanPlan, + fallbackScanId, + fallbackRequiredSchema)) + } + + private def descriptorFor( + scanExec: BatchScanExec, + nativeScanPlan: Option[Any]): LanceNativeScanDescriptor = { + val fallbackScanId = scanKey(scanExec) + val fallbackRequiredSchema = scanExec.scan.readSchema() + nativeScanPlan + .map(descriptorFromNativePlan(_, fallbackScanId, fallbackRequiredSchema)) + .getOrElse(fallbackDescriptor(fallbackScanId, fallbackRequiredSchema)) + } + + private def descriptorFromNativePlan( + nativeScanPlan: Any, + fallbackScanId: String, + fallbackRequiredSchema: StructType): LanceNativeScanDescriptor = { + val requiredSchema = + structTypeFromJson( + requireString(invokeRequired(nativeScanPlan, "getSparkReadSchemaJson")), + "getSparkReadSchemaJson") + val projectedSchema = + structTypeFromJson( + requireString(invokeRequired(nativeScanPlan, "getProjectedReadSchemaJson")), + "getProjectedReadSchemaJson") - val partitionScan = OperatorOuterClass.LanceScan + LanceNativeScanDescriptor( + descriptorVersion = toUInt32( + invokeRequired(nativeScanPlan, "getDescriptorVersion"), + "getDescriptorVersion"), + scanId = nonEmptyString( + invokeRequired(nativeScanPlan, "getScanId"), + fallbackScanId), + datasetUri = requireString(invokeRequired(nativeScanPlan, "getDatasetUri")), + resolvedVersion = toLong(invokeRequired(nativeScanPlan, "getResolvedVersion")), + storageOptions = toStringMap(invokeRequired(nativeScanPlan, "getStorageOptions")), + requiredSchema = requiredSchema, + projectedSchema = projectedSchema, + filterSql = optionalString(nativeScanPlan, "hasPushedFilterSql", "getPushedFilterSql"), + limit = optionalLong(nativeScanPlan, "hasLimit", "getLimit"), + offset = optionalLong(nativeScanPlan, "hasOffset", "getOffset"), + batchSize = toUInt32(invokeRequired(nativeScanPlan, "getBatchSize"), "getBatchSize"), + nativeScanPlanClass = nativeScanPlan.getClass.getName, + splits = toSeq(invokeRequired(nativeScanPlan, "getSplits")).map(splitFromNativeSplit)) + } + + private def fallbackDescriptor( + scanId: String, + requiredSchema: StructType): LanceNativeScanDescriptor = + LanceNativeScanDescriptor( + descriptorVersion = 0, + scanId = scanId, + datasetUri = "", + resolvedVersion = 0L, + storageOptions = Map.empty, + requiredSchema = requiredSchema, + projectedSchema = requiredSchema, + filterSql = None, + limit = None, + offset = None, + batchSize = 0, + nativeScanPlanClass = "", + splits = Seq(LanceNativeScanSplitDescriptor(0, Nil))) + + private def commonFromDescriptor( + descriptor: LanceNativeScanDescriptor): OperatorOuterClass.LanceScanCommon = { + val commonBuilder = OperatorOuterClass.LanceScanCommon .newBuilder() - .setPartition(partition) - .build() + .setScanId(descriptor.scanId) + .setNativeScanPlanClass(descriptor.nativeScanPlanClass) + .setDatasetUri(descriptor.datasetUri) + .setResolvedVersion(descriptor.resolvedVersion) + .putAllStorageOptions(descriptor.storageOptions.asJava) + .addAllRequiredSchema(schema2Proto(descriptor.requiredSchema.fields).toSeq.asJava) + .addAllProjectedSchema(schema2Proto(descriptor.projectedSchema.fields).toSeq.asJava) + .setBatchSize(descriptor.batchSize) + .setDescriptorVersion(descriptor.descriptorVersion) - (common.toByteArray, Array(partitionScan.toByteArray)) + descriptor.filterSql.foreach(commonBuilder.setFilterSql) + descriptor.limit.foreach(commonBuilder.setLimit) + descriptor.offset.foreach(commonBuilder.setOffset) + commonBuilder.build() } + private def splitFromNativeSplit(nativeSplit: Any): LanceNativeScanSplitDescriptor = + LanceNativeScanSplitDescriptor( + partitionIndex = toUInt32(invokeRequired(nativeSplit, "getSplitIndex"), "getSplitIndex"), + fragmentIds = toSeq(invokeRequired(nativeSplit, "getFragmentIds")) + .map(toUInt32(_, "getFragmentIds"))) + + private def structTypeFromJson(json: String, methodName: String): StructType = + try { + DataType.fromJson(json) match { + case schema: StructType => schema + case other => + throw new IllegalArgumentException( + s"expected StructType JSON but got ${other.typeName}") + } + } catch { + case NonFatal(e) => + throw new IllegalArgumentException( + s"Native Lance scan descriptor method $methodName returned invalid Spark schema JSON", + e) + } + + private def optionalString(target: Any, hasMethod: String, valueMethod: String): Option[String] = + if (toBoolean(invokeRequired(target, hasMethod))) { + Some(requireString(invokeRequired(target, valueMethod))) + } else { + None + } + + private def optionalLong(target: Any, hasMethod: String, valueMethod: String): Option[Long] = + if (toBoolean(invokeRequired(target, hasMethod))) { + Some(toLong(invokeRequired(target, valueMethod))) + } else { + None + } + + private def invokeRequired(target: Any, methodName: String): Any = { + require(target != null, s"Native Lance scan descriptor target is null for $methodName") + try { + findNoArgMethod(target.getClass, methodName) + .getOrElse { + throw new NoSuchMethodException(s"${target.getClass.getName}.$methodName()") + } + .invoke(target.asInstanceOf[AnyRef]) + } catch { + case e: InvocationTargetException if e.getCause != null => + throw e.getCause + case NonFatal(e) => + throw new IllegalArgumentException( + s"Unable to read native Lance scan descriptor method $methodName", + e) + } + } + + private def findNoArgMethod( + clazz: Class[_], + methodName: String): Option[java.lang.reflect.Method] = { + var current = clazz + while (current != null) { + try { + val method = current.getDeclaredMethod(methodName) + method.setAccessible(true) + return Some(method) + } catch { + case _: NoSuchMethodException => + current = current.getSuperclass + } + } + None + } + + private def toSeq(value: Any): Seq[Any] = value match { + case null => Seq.empty + case values: java.lang.Iterable[_] => values.asScala.toSeq + case values: Iterable[_] => values.toSeq + case values: Array[_] => values.toSeq + case other => + throw new IllegalArgumentException( + s"Expected a collection in native Lance scan descriptor, got ${other.getClass.getName}") + } + + private def toStringMap(value: Any): Map[String, String] = value match { + case null => Map.empty + case values: java.util.Map[_, _] => + values.asScala.map { case (key, value) => key.toString -> value.toString }.toMap + case values: collection.Map[_, _] => + values.map { case (key, value) => key.toString -> value.toString }.toMap + case other => + throw new IllegalArgumentException( + s"Expected a map in native Lance scan descriptor, got ${other.getClass.getName}") + } + + private def toBoolean(value: Any): Boolean = value match { + case value: java.lang.Boolean => value.booleanValue() + case value: Boolean => value + case other => + throw new IllegalArgumentException( + s"Expected boolean in native Lance scan descriptor, got ${typeName(other)}") + } + + private def toLong(value: Any): Long = value match { + case value: java.lang.Number => value.longValue() + case value: String => value.toLong + case other => + throw new IllegalArgumentException( + s"Expected integer in native Lance scan descriptor, got ${typeName(other)}") + } + + private def toUInt32(value: Any, methodName: String): Int = { + val longValue = toLong(value) + if (longValue < 0 || longValue > 0xffffffffL) { + throw new IllegalArgumentException( + s"Native Lance scan descriptor method $methodName returned out-of-range uint32 " + + s"value $longValue") + } + longValue.toInt + } + + private def requireString(value: Any): String = value match { + case null => "" + case value: String => value + case other => other.toString + } + + private def nonEmptyString(value: Any, fallback: String): String = { + val stringValue = requireString(value) + if (stringValue.nonEmpty) stringValue else fallback + } + + private def typeName(value: Any): String = + Option(value).map(_.getClass.getName).getOrElse("null") + private def scanKey(scanExec: BatchScanExec): String = s"lance_${scanExec.id}_${scanExec.scan.hashCode()}" } diff --git a/spark/src/contrib-lance/scala/org/apache/spark/sql/comet/CometLanceNativeScanExec.scala b/spark/src/contrib-lance/scala/org/apache/spark/sql/comet/CometLanceNativeScanExec.scala index da82e06e421..5843e925364 100644 --- a/spark/src/contrib-lance/scala/org/apache/spark/sql/comet/CometLanceNativeScanExec.scala +++ b/spark/src/contrib-lance/scala/org/apache/spark/sql/comet/CometLanceNativeScanExec.scala @@ -21,28 +21,28 @@ package org.apache.spark.sql.comet import scala.jdk.CollectionConverters._ +import org.apache.spark.{Partition, TaskContext} import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, SortOrder} import org.apache.spark.sql.catalyst.plans.QueryPlan import org.apache.spark.sql.catalyst.plans.physical.{Partitioning, UnknownPartitioning} import org.apache.spark.sql.execution.datasources.v2.BatchScanExec -import org.apache.spark.sql.types.StructType import org.apache.spark.sql.vectorized.ColumnarBatch import com.google.common.base.Objects import org.apache.comet.serde.OperatorOuterClass.Operator import org.apache.comet.serde.operator.CometLanceNativeScan +import org.apache.comet.serde.operator.CometLanceNativeScan.LanceNativeScanDescriptor case class CometLanceNativeScanExec( override val nativeOp: Operator, override val output: Seq[Attribute], runtimeFilters: Seq[Expression], - requiredSchema: StructType, @transient originalPlan: BatchScanExec, override val serializedPlanOpt: SerializedPlan, override val sourceKey: String, - nativeScanPlanClassName: String) + lanceDescriptor: LanceNativeScanDescriptor) extends CometLeafExec with CometLanceNativeScanLike { @@ -51,7 +51,7 @@ case class CometLanceNativeScanExec( override val nodeName: String = "CometLanceNativeScan" @transient private lazy val serializedPartitionData: (Array[Byte], Array[Array[Byte]]) = - CometLanceNativeScan.serializePartitions(sourceKey, requiredSchema, nativeScanPlanClassName) + CometLanceNativeScan.serializePartitions(lanceDescriptor) override def commonData: Array[Byte] = serializedPartitionData._1 @@ -63,8 +63,24 @@ case class CometLanceNativeScanExec( override lazy val outputOrdering: Seq[SortOrder] = Nil override def doExecuteColumnar(): RDD[ColumnarBatch] = { - throw new UnsupportedOperationException( - "Native Lance scan execution is not implemented yet") + val nativeMetrics = CometMetricNode.fromCometPlan(this) + val serializedPlan = CometExec.serializeNativePlan(nativeOp) + new CometExecRDD( + sparkContext, + inputRDDs = Seq.empty, + commonByKey = Map(sourceKey -> commonData), + perPartitionByKey = Map(sourceKey -> perPartitionData), + serializedPlan = serializedPlan, + defaultNumPartitions = perPartitionData.length, + numOutputCols = output.length, + nativeMetrics = nativeMetrics, + subqueries = Seq.empty) { + override def compute(split: Partition, context: TaskContext): Iterator[ColumnarBatch] = { + val res = super.compute(split, context) + Option(context).foreach(nativeMetrics.reportScanInputMetrics) + res + } + } } override def convertBlock(): CometLanceNativeScanExec = { @@ -78,11 +94,10 @@ case class CometLanceNativeScanExec( nativeOp, output, runtimeFilters, - requiredSchema, originalPlan, newSerializedPlan, sourceKey, - nativeScanPlanClassName) + lanceDescriptor) } override protected def doCanonicalize(): CometLanceNativeScanExec = { @@ -92,15 +107,14 @@ case class CometLanceNativeScanExec( QueryPlan.normalizePredicates( CometScanUtils.filterUnusedDynamicPruningExpressions(runtimeFilters), output), - requiredSchema, null, SerializedPlan(None), sourceKey, - nativeScanPlanClassName) + lanceDescriptor) } override def stringArgs: Iterator[Any] = - Iterator(output, s"$sourceKey, nativeScanPlan=$nativeScanPlanClassName") + Iterator(output, s"$sourceKey, nativeScanPlan=${lanceDescriptor.nativeScanPlanClass}") override def equals(obj: Any): Boolean = obj match { case other: CometLanceNativeScanExec => diff --git a/spark/src/test/scala/org/apache/comet/rules/CometScanRuleSuite.scala b/spark/src/test/scala/org/apache/comet/rules/CometScanRuleSuite.scala index 521afa44a9d..15a9de3dde3 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometScanRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometScanRuleSuite.scala @@ -19,7 +19,11 @@ package org.apache.comet.rules +import java.util.{Arrays, LinkedHashMap} + +import scala.jdk.CollectionConverters._ import scala.util.Random +import scala.util.Try import org.apache.spark.sql._ import org.apache.spark.sql.comet._ @@ -28,6 +32,7 @@ import org.apache.spark.sql.execution.adaptive.QueryStageExec import org.apache.spark.sql.types.{DataTypes, StructField, StructType} import org.apache.comet.CometConf +import org.apache.comet.serde.OperatorOuterClass import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator} /** @@ -185,4 +190,107 @@ class CometScanRuleSuite extends CometTestBase { assert(!CometConf.COMET_LANCE_NATIVE_ENABLED.get()) } + test("Lance native scan serde reflects descriptor common fields and split fragments") { + val serde = loadContribLanceSerde.getOrElse { + cancel("contrib-lance profile is not enabled") + } + + val requiredSchema = StructType( + Seq( + StructField("id", DataTypes.IntegerType, nullable = false), + StructField("name", DataTypes.StringType, nullable = true))) + val projectedSchema = StructType(Seq(StructField("id", DataTypes.IntegerType, false))) + val descriptor = new FakeLanceNativeScanPlan(requiredSchema, projectedSchema) + + val (common, partitions) = + serializeFakeLanceDescriptor(serde, descriptor, "fallback-scan", requiredSchema) + + assert(common.getScanId == "scan-123") + assert(common.getDatasetUri == "s3://bucket/table.lance") + assert(common.getResolvedVersion == 42L) + assert(common.getDescriptorVersion == 1) + assert(common.getBatchSize == 4096) + assert(common.getNativeScanPlanClass.contains("FakeLanceNativeScanPlan")) + assert(common.getStorageOptionsMap.get("region") == "us-west-2") + assert(common.getStorageOptionsMap.get("endpoint") == "http://127.0.0.1:9000") + assert(common.getRequiredSchemaList.asScala.map(_.getName) == Seq("id", "name")) + assert(common.getProjectedSchemaList.asScala.map(_.getName) == Seq("id")) + assert(common.hasFilterSql) + assert(common.getFilterSql == "id > 10") + assert(common.hasLimit) + assert(common.getLimit == 100L) + assert(common.hasOffset) + assert(common.getOffset == 5L) + + assert(partitions.length == 2) + assert(partitions(0).getPartition.getPartitionIndex == 0) + assert(partitions(0).getPartition.getFragmentIdsList.asScala.map(_.intValue()) == Seq(7, 8)) + assert(partitions(1).getPartition.getPartitionIndex == 1) + assert(partitions(1).getPartition.getFragmentIdsList.asScala.map(_.intValue()) == Seq(9)) + } + + private def loadContribLanceSerde: Option[AnyRef] = + Try { + Class + .forName("org.apache.comet.serde.operator.CometLanceNativeScan$") + .getField("MODULE$") + .get(null) + .asInstanceOf[AnyRef] + }.toOption + + private def serializeFakeLanceDescriptor( + serde: AnyRef, + descriptor: AnyRef, + fallbackScanId: String, + fallbackRequiredSchema: StructType) + : (OperatorOuterClass.LanceScanCommon, Array[OperatorOuterClass.LanceScan]) = { + val method = serde.getClass.getMethods + .find(method => + method.getName == "serializeNativePlan" && method.getParameterTypes.length == 3) + .getOrElse { + throw new AssertionError("CometLanceNativeScan.serializeNativePlan was not found") + } + + val serialized = method + .invoke(serde, descriptor, fallbackScanId, fallbackRequiredSchema) + .asInstanceOf[Product] + val commonBytes = serialized.productElement(0).asInstanceOf[Array[Byte]] + val partitionBytes = serialized.productElement(1).asInstanceOf[Array[Array[Byte]]] + + ( + OperatorOuterClass.LanceScanCommon.parseFrom(commonBytes), + partitionBytes.map(OperatorOuterClass.LanceScan.parseFrom)) + } + + private class FakeLanceNativeScanPlan( + requiredSchema: StructType, + projectedSchema: StructType) { + private val storageOptions = new LinkedHashMap[String, String]() + storageOptions.put("region", "us-west-2") + storageOptions.put("endpoint", "http://127.0.0.1:9000") + + def getDescriptorVersion(): Int = 1 + def getScanId(): String = "scan-123" + def getDatasetUri(): String = "s3://bucket/table.lance" + def getResolvedVersion(): Long = 42L + def getSparkReadSchemaJson(): String = requiredSchema.json + def getProjectedReadSchemaJson(): String = projectedSchema.json + def hasPushedFilterSql(): Boolean = true + def getPushedFilterSql(): String = "id > 10" + def hasLimit(): Boolean = true + def getLimit(): Long = 100L + def hasOffset(): Boolean = true + def getOffset(): Long = 5L + def getBatchSize(): Int = 4096 + def getStorageOptions(): java.util.Map[String, String] = storageOptions + def getSplits(): java.util.List[FakeLanceNativeScanSplit] = + Arrays.asList( + new FakeLanceNativeScanSplit(0, Arrays.asList(Int.box(7), Int.box(8))), + new FakeLanceNativeScanSplit(1, Arrays.asList(Int.box(9)))) + } + + private class FakeLanceNativeScanSplit(splitIndex: Int, fragmentIds: java.util.List[Integer]) { + def getSplitIndex(): Int = splitIndex + def getFragmentIds(): java.util.List[Integer] = fragmentIds + } } From 87156c42e6e81ed437fb2d49643bf7a406635be0 Mon Sep 17 00:00:00 2001 From: xuanyili Date: Fri, 12 Jun 2026 07:58:47 +0000 Subject: [PATCH 4/6] fix: unwrap Lance native scan descriptor --- .../apache/comet/lance/LanceIntegration.scala | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/spark/src/main/scala/org/apache/comet/lance/LanceIntegration.scala b/spark/src/main/scala/org/apache/comet/lance/LanceIntegration.scala index 348cb799f9c..4c37aefa043 100644 --- a/spark/src/main/scala/org/apache/comet/lance/LanceIntegration.scala +++ b/spark/src/main/scala/org/apache/comet/lance/LanceIntegration.scala @@ -20,6 +20,7 @@ package org.apache.comet.lance import java.lang.reflect.InvocationTargetException +import java.util.{Optional => JOptional} import scala.util.control.NonFatal @@ -104,6 +105,13 @@ object LanceIntegration extends Logging { None } } catch { + case e: InvocationTargetException => + val cause = Option(e.getCause).getOrElse(e) + logWarning( + "Native Lance scan disabled because contrib-lance threw during reflection: " + + s"${cause.getClass.getName}: ${cause.getMessage}", + cause) + None case NonFatal(e) => logWarning(s"Native Lance scan disabled by contrib-lance reflection failure: $e") None @@ -114,7 +122,7 @@ object LanceIntegration extends Logging { try { findNoArgMethod(scan.getClass, NativeScanPlanMethod) .flatMap { method => - Option(method.invoke(scan)) + optionalResult(method.invoke(scan)) } } catch { case e: InvocationTargetException => @@ -128,6 +136,14 @@ object LanceIntegration extends Logging { } } + private def optionalResult(value: Any): Option[Any] = value match { + case null => None + case option: Option[_] => option + case option: JOptional[_] if option.isPresent => Some(option.get) + case _: JOptional[_] => None + case other => Some(other) + } + private def findNoArgMethod( clazz: Class[_], methodName: String): Option[java.lang.reflect.Method] = { From 0cc90357b365d97708de7ed429e44ed560b0b9fb Mon Sep 17 00:00:00 2001 From: xuanyili Date: Fri, 12 Jun 2026 08:19:54 +0000 Subject: [PATCH 5/6] fix: adapt Lance fallback reasons --- .../org/apache/comet/lance/CometLanceSupport.scala | 4 ++-- .../org/apache/comet/lance/LanceIntegration.scala | 10 +++++----- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/spark/src/contrib-lance/scala/org/apache/comet/lance/CometLanceSupport.scala b/spark/src/contrib-lance/scala/org/apache/comet/lance/CometLanceSupport.scala index d02f73dec5d..69d2ed2a440 100644 --- a/spark/src/contrib-lance/scala/org/apache/comet/lance/CometLanceSupport.scala +++ b/spark/src/contrib-lance/scala/org/apache/comet/lance/CometLanceSupport.scala @@ -25,7 +25,7 @@ import org.apache.spark.sql.comet.CometBatchScanExec import org.apache.spark.sql.execution.SparkPlan import org.apache.spark.sql.execution.datasources.v2.BatchScanExec -import org.apache.comet.CometSparkSessionExtensions.withInfos +import org.apache.comet.CometSparkSessionExtensions.withFallbackReasons import org.apache.comet.serde.OperatorOuterClass.Operator import org.apache.comet.serde.operator.CometLanceNativeScan @@ -38,7 +38,7 @@ object CometLanceSupport { if (!schemaSupported) { fallbackReasons += s"Schema ${scanExec.scan.readSchema()} is not supported" - withInfos(scanExec, fallbackReasons.toSet) + withFallbackReasons(scanExec, fallbackReasons.toSet) None } else { val builder = Operator.newBuilder().setPlanId(scanExec.id) diff --git a/spark/src/main/scala/org/apache/comet/lance/LanceIntegration.scala b/spark/src/main/scala/org/apache/comet/lance/LanceIntegration.scala index 4c37aefa043..5424bb5f6d9 100644 --- a/spark/src/main/scala/org/apache/comet/lance/LanceIntegration.scala +++ b/spark/src/main/scala/org/apache/comet/lance/LanceIntegration.scala @@ -29,7 +29,7 @@ import org.apache.spark.sql.execution.SparkPlan import org.apache.spark.sql.execution.datasources.v2.BatchScanExec import org.apache.comet.CometConf -import org.apache.comet.CometSparkSessionExtensions.withInfo +import org.apache.comet.CometSparkSessionExtensions.withFallbackReason /** * Reflection-only bridge for optional Lance Spark integration. @@ -60,7 +60,7 @@ object LanceIntegration extends Logging { def tryCreateNativeScan(scanExec: BatchScanExec): Option[SparkPlan] = { if (!CometConf.COMET_LANCE_NATIVE_ENABLED.get(scanExec.conf)) { - withInfo( + withFallbackReason( scanExec, s"Native Lance scan disabled because ${CometConf.COMET_LANCE_NATIVE_ENABLED.key} " + "is not enabled") @@ -68,7 +68,7 @@ object LanceIntegration extends Logging { } if (!CometConf.COMET_EXEC_ENABLED.get(scanExec.conf)) { - withInfo( + withFallbackReason( scanExec, s"Native Lance scan disabled because ${CometConf.COMET_EXEC_ENABLED.key} is not enabled") return None @@ -77,7 +77,7 @@ object LanceIntegration extends Logging { val nativePlan = nativeScanPlan(scanExec.scan) match { case Some(plan) => plan case None => - withInfo( + withFallbackReason( scanExec, s"Native Lance scan disabled because $LanceScanClassName.$NativeScanPlanMethod() " + "is not available") @@ -87,7 +87,7 @@ object LanceIntegration extends Logging { val support = loadContribSupport match { case Some(module) => module case None => - withInfo( + withFallbackReason( scanExec, "Native Lance scan disabled because the contrib-lance build profile is not present") return None From 2a0ba55dff4d1bcb963bf8a5a2e0cc14969883f0 Mon Sep 17 00:00:00 2001 From: xuanyili Date: Tue, 14 Jul 2026 02:46:27 +0000 Subject: [PATCH 6/6] refactor: move Lance scan hooks to contrib SPI --- spark/pom.xml | 14 + .../org.apache.comet.rules.CometScanContrib | 1 + ...rg.apache.spark.sql.comet.PlanDataInjector | 1 + .../comet/lance/CometLanceSupport.scala | 37 ++- .../comet/lance/LanceScanRuleExtension.scala | 96 ++++++ .../serde/operator/CometLanceNativeScan.scala | 313 +++++------------- .../sql/comet/CometLanceNativeScanExec.scala | 2 +- .../sql/comet/LancePlanDataInjector.scala | 57 ++++ .../apache/comet/lance/LanceIntegration.scala | 194 ----------- .../sql/comet/CometLanceNativeScanLike.scala | 28 -- .../comet/rules/CometScanRuleSuite.scala | 17 +- 11 files changed, 292 insertions(+), 468 deletions(-) create mode 100644 spark/src/contrib-lance/resources/META-INF/services/org.apache.comet.rules.CometScanContrib create mode 100644 spark/src/contrib-lance/resources/META-INF/services/org.apache.spark.sql.comet.PlanDataInjector create mode 100644 spark/src/contrib-lance/scala/org/apache/comet/lance/LanceScanRuleExtension.scala create mode 100644 spark/src/contrib-lance/scala/org/apache/spark/sql/comet/LancePlanDataInjector.scala delete mode 100644 spark/src/main/scala/org/apache/comet/lance/LanceIntegration.scala delete mode 100644 spark/src/main/scala/org/apache/spark/sql/comet/CometLanceNativeScanLike.scala diff --git a/spark/pom.xml b/spark/pom.xml index 19341cd5127..4a9c2b76bfd 100644 --- a/spark/pom.xml +++ b/spark/pom.xml @@ -525,6 +525,20 @@ under the License. + + add-contrib-lance-resources + generate-resources + + add-resource + + + + + src/contrib-lance/resources + + + + diff --git a/spark/src/contrib-lance/resources/META-INF/services/org.apache.comet.rules.CometScanContrib b/spark/src/contrib-lance/resources/META-INF/services/org.apache.comet.rules.CometScanContrib new file mode 100644 index 00000000000..4f7e48a550b --- /dev/null +++ b/spark/src/contrib-lance/resources/META-INF/services/org.apache.comet.rules.CometScanContrib @@ -0,0 +1 @@ +org.apache.comet.lance.LanceScanRuleExtension diff --git a/spark/src/contrib-lance/resources/META-INF/services/org.apache.spark.sql.comet.PlanDataInjector b/spark/src/contrib-lance/resources/META-INF/services/org.apache.spark.sql.comet.PlanDataInjector new file mode 100644 index 00000000000..45ba51e9074 --- /dev/null +++ b/spark/src/contrib-lance/resources/META-INF/services/org.apache.spark.sql.comet.PlanDataInjector @@ -0,0 +1 @@ +org.apache.spark.sql.comet.LancePlanDataInjector diff --git a/spark/src/contrib-lance/scala/org/apache/comet/lance/CometLanceSupport.scala b/spark/src/contrib-lance/scala/org/apache/comet/lance/CometLanceSupport.scala index 69d2ed2a440..0fc46b5a844 100644 --- a/spark/src/contrib-lance/scala/org/apache/comet/lance/CometLanceSupport.scala +++ b/spark/src/contrib-lance/scala/org/apache/comet/lance/CometLanceSupport.scala @@ -21,12 +21,19 @@ package org.apache.comet.lance import scala.collection.mutable.ListBuffer +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, SortOrder} import org.apache.spark.sql.comet.CometBatchScanExec +import org.apache.spark.sql.catalyst.plans.physical.Partitioning import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.LeafExecNode import org.apache.spark.sql.execution.datasources.v2.BatchScanExec +import org.apache.spark.sql.vectorized.ColumnarBatch import org.apache.comet.CometSparkSessionExtensions.withFallbackReasons -import org.apache.comet.serde.OperatorOuterClass.Operator +import org.apache.comet.rules.CometContribScanMarker +import org.apache.comet.serde.CometOperatorSerde import org.apache.comet.serde.operator.CometLanceNativeScan object CometLanceSupport { @@ -41,12 +48,28 @@ object CometLanceSupport { withFallbackReasons(scanExec, fallbackReasons.toSet) None } else { - val builder = Operator.newBuilder().setPlanId(scanExec.id) - CometLanceNativeScan - .convert(scanExec, builder, Option(nativeScanPlan)) - .map { nativeOp => - CometLanceNativeScan.createExec(nativeOp, scanExec, Option(nativeScanPlan)) - } + val marker = LanceScanExec(scanExec, nativeScanPlan) + scanExec.logicalLink.foreach(marker.setLogicalLink) + Some(marker) } } } + +case class LanceScanExec(originalPlan: BatchScanExec, nativeScanPlan: Object) + extends LeafExecNode + with CometContribScanMarker { + + override def scanHandler: CometOperatorSerde[_ <: SparkPlan] = CometLanceNativeScan + + override def output: Seq[Attribute] = originalPlan.output + + override def outputPartitioning: Partitioning = originalPlan.outputPartitioning + + override def outputOrdering: Seq[SortOrder] = originalPlan.outputOrdering + + override def supportsColumnar: Boolean = originalPlan.supportsColumnar + + override protected def doExecute(): RDD[InternalRow] = originalPlan.execute() + + override protected def doExecuteColumnar(): RDD[ColumnarBatch] = originalPlan.executeColumnar() +} diff --git a/spark/src/contrib-lance/scala/org/apache/comet/lance/LanceScanRuleExtension.scala b/spark/src/contrib-lance/scala/org/apache/comet/lance/LanceScanRuleExtension.scala new file mode 100644 index 00000000000..693fb487621 --- /dev/null +++ b/spark/src/contrib-lance/scala/org/apache/comet/lance/LanceScanRuleExtension.scala @@ -0,0 +1,96 @@ +/* + * 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 org.apache.comet.lance + +import java.util.{Optional => JOptional} + +import scala.util.control.NonFatal + +import org.apache.spark.internal.Logging +import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.datasources.v2.BatchScanExec + +import org.apache.comet.CometConf +import org.apache.comet.CometSparkSessionExtensions.withFallbackReason +import org.apache.comet.rules.CometScanContrib + +/** + * `CometScanContrib` implementation for Lance V2 scans. + * + * The class is loaded through `ServiceLoader` from the optional contrib-lance profile. It keeps + * Lance Spark references reflective so building this contrib does not require a lance-spark + * compile-time dependency. + */ +class LanceScanRuleExtension extends CometScanContrib with Logging { + + private val LanceScanClassName = "org.lance.spark.read.LanceScan" + private val NativeScanPlanMethod = "nativeScanPlan" + + override def tryTransformV2(scanExec: BatchScanExec): Option[SparkPlan] = { + if (scanExec.scan.getClass.getName != LanceScanClassName) { + return None + } + + Some(tryCreateNativeScan(scanExec).getOrElse(scanExec)) + } + + private def nativeScanPlan(scan: AnyRef): Option[AnyRef] = { + try { + val result = scan.getClass + .getMethod(NativeScanPlanMethod) + .invoke(scan) + .asInstanceOf[JOptional[_]] + if (result.isPresent) Some(result.get().asInstanceOf[AnyRef]) else None + } catch { + case NonFatal(e) => + logWarning(s"Native Lance scan disabled because $NativeScanPlanMethod() failed", e) + None + } + } + + private def tryCreateNativeScan(scanExec: BatchScanExec): Option[SparkPlan] = { + if (!CometConf.COMET_LANCE_NATIVE_ENABLED.get(scanExec.conf)) { + withFallbackReason( + scanExec, + s"Native Lance scan disabled because ${CometConf.COMET_LANCE_NATIVE_ENABLED.key} " + + "is not enabled") + return None + } + + if (!CometConf.COMET_EXEC_ENABLED.get(scanExec.conf)) { + withFallbackReason( + scanExec, + s"Native Lance scan disabled because ${CometConf.COMET_EXEC_ENABLED.key} is not enabled") + return None + } + + val nativePlan = nativeScanPlan(scanExec.scan.asInstanceOf[AnyRef]) match { + case Some(plan) => plan + case None => + withFallbackReason( + scanExec, + s"Native Lance scan disabled because $LanceScanClassName.$NativeScanPlanMethod() " + + "is not available") + return None + } + + CometLanceSupport.tryTransform(scanExec, nativePlan) + } +} diff --git a/spark/src/contrib-lance/scala/org/apache/comet/serde/operator/CometLanceNativeScan.scala b/spark/src/contrib-lance/scala/org/apache/comet/serde/operator/CometLanceNativeScan.scala index e20ecf1df89..f0bcb37c477 100644 --- a/spark/src/contrib-lance/scala/org/apache/comet/serde/operator/CometLanceNativeScan.scala +++ b/spark/src/contrib-lance/scala/org/apache/comet/serde/operator/CometLanceNativeScan.scala @@ -19,21 +19,19 @@ package org.apache.comet.serde.operator -import java.lang.reflect.InvocationTargetException - import scala.jdk.CollectionConverters._ -import scala.util.control.NonFatal -import org.apache.spark.internal.Logging import org.apache.spark.sql.comet.{CometLanceNativeScanExec, CometNativeExec, SerializedPlan} -import org.apache.spark.sql.execution.datasources.v2.BatchScanExec import org.apache.spark.sql.types.{DataType, StructType} import org.apache.comet.{CometConf, ConfigEntry} +import org.apache.comet.lance.LanceScanExec import org.apache.comet.serde.{CometOperatorSerde, Compatible, OperatorOuterClass, SupportLevel} import org.apache.comet.serde.OperatorOuterClass.Operator -object CometLanceNativeScan extends CometOperatorSerde[BatchScanExec] with Logging { +object CometLanceNativeScan extends CometOperatorSerde[LanceScanExec] { + + val LanceScanTypeUrl = "type.googleapis.com/comet.contrib.lance.LanceScan" case class LanceNativeScanSplitDescriptor(partitionIndex: Int, fragmentIds: Seq[Int]) @@ -55,45 +53,38 @@ object CometLanceNativeScan extends CometOperatorSerde[BatchScanExec] with Loggi override def enabledConfig: Option[ConfigEntry[Boolean]] = Some(CometConf.COMET_LANCE_NATIVE_ENABLED) - override def getSupportLevel(operator: BatchScanExec): SupportLevel = Compatible() + override def getSupportLevel(operator: LanceScanExec): SupportLevel = Compatible() override def convert( - scanExec: BatchScanExec, - builder: Operator.Builder, - childOp: Operator*): Option[Operator] = - convert(scanExec, builder, None) - - def convert( - scanExec: BatchScanExec, + scanExec: LanceScanExec, builder: Operator.Builder, - nativeScanPlan: Option[Any]): Option[Operator] = { - val descriptor = descriptorFor(scanExec, nativeScanPlan) - - val lanceScanBuilder = OperatorOuterClass.LanceScan + childOp: Operator*): Option[Operator] = { + val descriptor = descriptorFromNativePlan(scanExec.nativeScanPlan) + val lanceScan = OperatorOuterClass.LanceScan .newBuilder() .setCommon(commonFromDescriptor(descriptor)) + .build() + val contribScan = OperatorOuterClass.ContribScan + .newBuilder() + .setTypeUrl(LanceScanTypeUrl) + .setValue(lanceScan.toByteString) builder.clearChildren() - Some(builder.setLanceScan(lanceScanBuilder).build()) + Some(builder.setContribScan(contribScan).build()) } - override def createExec(nativeOp: Operator, op: BatchScanExec): CometNativeExec = - createExec(nativeOp, op, None) - - def createExec( - nativeOp: Operator, - op: BatchScanExec, - nativeScanPlan: Option[Any]): CometNativeExec = { - val descriptor = descriptorFor(op, nativeScanPlan) + override def createExec(nativeOp: Operator, op: LanceScanExec): CometNativeExec = { + val descriptor = descriptorFromNativePlan(op.nativeScanPlan) + val scan = op.originalPlan val exec = CometLanceNativeScanExec( nativeOp, - op.output, - op.runtimeFilters, - op, + scan.output, + scan.runtimeFilters, + scan, SerializedPlan(None), descriptor.scanId, descriptor) - op.logicalLink.foreach(exec.setLogicalLink) + scan.logicalLink.foreach(exec.setLogicalLink) exec } @@ -115,76 +106,80 @@ object CometLanceNativeScan extends CometOperatorSerde[BatchScanExec] with Loggi }.toArray) private[comet] def serializeNativePlan( - nativeScanPlan: Any, - fallbackScanId: String, - fallbackRequiredSchema: StructType): (Array[Byte], Array[Array[Byte]]) = { - serializePartitions(descriptorFromNativePlan( - nativeScanPlan, - fallbackScanId, - fallbackRequiredSchema)) - } - - private def descriptorFor( - scanExec: BatchScanExec, - nativeScanPlan: Option[Any]): LanceNativeScanDescriptor = { - val fallbackScanId = scanKey(scanExec) - val fallbackRequiredSchema = scanExec.scan.readSchema() - nativeScanPlan - .map(descriptorFromNativePlan(_, fallbackScanId, fallbackRequiredSchema)) - .getOrElse(fallbackDescriptor(fallbackScanId, fallbackRequiredSchema)) - } + nativeScanPlan: Object): (Array[Byte], Array[Array[Byte]]) = + serializePartitions(descriptorFromNativePlan(nativeScanPlan)) private def descriptorFromNativePlan( - nativeScanPlan: Any, - fallbackScanId: String, - fallbackRequiredSchema: StructType): LanceNativeScanDescriptor = { + nativeScanPlan: Object): LanceNativeScanDescriptor = { + def invoke(methodName: String): AnyRef = + nativeScanPlan.getClass.getMethod(methodName).invoke(nativeScanPlan) + + def int(methodName: String): Int = + invoke(methodName).asInstanceOf[java.lang.Number].intValue() + + def long(methodName: String): Long = + invoke(methodName).asInstanceOf[java.lang.Number].longValue() + + def optionalString(hasMethod: String, valueMethod: String): Option[String] = + if (invoke(hasMethod).asInstanceOf[java.lang.Boolean].booleanValue()) { + Some(invoke(valueMethod).asInstanceOf[String]) + } else None + + def optionalLong(hasMethod: String, valueMethod: String): Option[Long] = + if (invoke(hasMethod).asInstanceOf[java.lang.Boolean].booleanValue()) { + Some(long(valueMethod)) + } else None + val requiredSchema = - structTypeFromJson( - requireString(invokeRequired(nativeScanPlan, "getSparkReadSchemaJson")), - "getSparkReadSchemaJson") + DataType + .fromJson(invoke("getSparkReadSchemaJson").asInstanceOf[String]) + .asInstanceOf[StructType] val projectedSchema = - structTypeFromJson( - requireString(invokeRequired(nativeScanPlan, "getProjectedReadSchemaJson")), - "getProjectedReadSchemaJson") + DataType + .fromJson(invoke("getProjectedReadSchemaJson").asInstanceOf[String]) + .asInstanceOf[StructType] + val storageOptions = invoke("getStorageOptions") + .asInstanceOf[java.util.Map[String, String]] + .asScala + .toMap + val splits = invoke("getSplits") + .asInstanceOf[java.lang.Iterable[Object]] + .asScala + .map { split => + val splitClass = split.getClass + val fragmentIds = splitClass + .getMethod("getFragmentIds") + .invoke(split) + .asInstanceOf[java.lang.Iterable[java.lang.Number]] + .asScala + .map(_.intValue()) + .toSeq + LanceNativeScanSplitDescriptor( + splitClass + .getMethod("getSplitIndex") + .invoke(split) + .asInstanceOf[java.lang.Number] + .intValue(), + fragmentIds) + } + .toSeq LanceNativeScanDescriptor( - descriptorVersion = toUInt32( - invokeRequired(nativeScanPlan, "getDescriptorVersion"), - "getDescriptorVersion"), - scanId = nonEmptyString( - invokeRequired(nativeScanPlan, "getScanId"), - fallbackScanId), - datasetUri = requireString(invokeRequired(nativeScanPlan, "getDatasetUri")), - resolvedVersion = toLong(invokeRequired(nativeScanPlan, "getResolvedVersion")), - storageOptions = toStringMap(invokeRequired(nativeScanPlan, "getStorageOptions")), + descriptorVersion = int("getDescriptorVersion"), + scanId = invoke("getScanId").asInstanceOf[String], + datasetUri = invoke("getDatasetUri").asInstanceOf[String], + resolvedVersion = long("getResolvedVersion"), + storageOptions = storageOptions, requiredSchema = requiredSchema, projectedSchema = projectedSchema, - filterSql = optionalString(nativeScanPlan, "hasPushedFilterSql", "getPushedFilterSql"), - limit = optionalLong(nativeScanPlan, "hasLimit", "getLimit"), - offset = optionalLong(nativeScanPlan, "hasOffset", "getOffset"), - batchSize = toUInt32(invokeRequired(nativeScanPlan, "getBatchSize"), "getBatchSize"), + filterSql = optionalString("hasPushedFilterSql", "getPushedFilterSql"), + limit = optionalLong("hasLimit", "getLimit"), + offset = optionalLong("hasOffset", "getOffset"), + batchSize = int("getBatchSize"), nativeScanPlanClass = nativeScanPlan.getClass.getName, - splits = toSeq(invokeRequired(nativeScanPlan, "getSplits")).map(splitFromNativeSplit)) + splits = splits) } - private def fallbackDescriptor( - scanId: String, - requiredSchema: StructType): LanceNativeScanDescriptor = - LanceNativeScanDescriptor( - descriptorVersion = 0, - scanId = scanId, - datasetUri = "", - resolvedVersion = 0L, - storageOptions = Map.empty, - requiredSchema = requiredSchema, - projectedSchema = requiredSchema, - filterSql = None, - limit = None, - offset = None, - batchSize = 0, - nativeScanPlanClass = "", - splits = Seq(LanceNativeScanSplitDescriptor(0, Nil))) - private def commonFromDescriptor( descriptor: LanceNativeScanDescriptor): OperatorOuterClass.LanceScanCommon = { val commonBuilder = OperatorOuterClass.LanceScanCommon @@ -194,8 +189,8 @@ object CometLanceNativeScan extends CometOperatorSerde[BatchScanExec] with Loggi .setDatasetUri(descriptor.datasetUri) .setResolvedVersion(descriptor.resolvedVersion) .putAllStorageOptions(descriptor.storageOptions.asJava) - .addAllRequiredSchema(schema2Proto(descriptor.requiredSchema.fields).toSeq.asJava) - .addAllProjectedSchema(schema2Proto(descriptor.projectedSchema.fields).toSeq.asJava) + .addAllRequiredSchema(schema2Proto(descriptor.requiredSchema.fields.toIndexedSeq).asJava) + .addAllProjectedSchema(schema2Proto(descriptor.projectedSchema.fields.toIndexedSeq).asJava) .setBatchSize(descriptor.batchSize) .setDescriptorVersion(descriptor.descriptorVersion) @@ -204,138 +199,4 @@ object CometLanceNativeScan extends CometOperatorSerde[BatchScanExec] with Loggi descriptor.offset.foreach(commonBuilder.setOffset) commonBuilder.build() } - - private def splitFromNativeSplit(nativeSplit: Any): LanceNativeScanSplitDescriptor = - LanceNativeScanSplitDescriptor( - partitionIndex = toUInt32(invokeRequired(nativeSplit, "getSplitIndex"), "getSplitIndex"), - fragmentIds = toSeq(invokeRequired(nativeSplit, "getFragmentIds")) - .map(toUInt32(_, "getFragmentIds"))) - - private def structTypeFromJson(json: String, methodName: String): StructType = - try { - DataType.fromJson(json) match { - case schema: StructType => schema - case other => - throw new IllegalArgumentException( - s"expected StructType JSON but got ${other.typeName}") - } - } catch { - case NonFatal(e) => - throw new IllegalArgumentException( - s"Native Lance scan descriptor method $methodName returned invalid Spark schema JSON", - e) - } - - private def optionalString(target: Any, hasMethod: String, valueMethod: String): Option[String] = - if (toBoolean(invokeRequired(target, hasMethod))) { - Some(requireString(invokeRequired(target, valueMethod))) - } else { - None - } - - private def optionalLong(target: Any, hasMethod: String, valueMethod: String): Option[Long] = - if (toBoolean(invokeRequired(target, hasMethod))) { - Some(toLong(invokeRequired(target, valueMethod))) - } else { - None - } - - private def invokeRequired(target: Any, methodName: String): Any = { - require(target != null, s"Native Lance scan descriptor target is null for $methodName") - try { - findNoArgMethod(target.getClass, methodName) - .getOrElse { - throw new NoSuchMethodException(s"${target.getClass.getName}.$methodName()") - } - .invoke(target.asInstanceOf[AnyRef]) - } catch { - case e: InvocationTargetException if e.getCause != null => - throw e.getCause - case NonFatal(e) => - throw new IllegalArgumentException( - s"Unable to read native Lance scan descriptor method $methodName", - e) - } - } - - private def findNoArgMethod( - clazz: Class[_], - methodName: String): Option[java.lang.reflect.Method] = { - var current = clazz - while (current != null) { - try { - val method = current.getDeclaredMethod(methodName) - method.setAccessible(true) - return Some(method) - } catch { - case _: NoSuchMethodException => - current = current.getSuperclass - } - } - None - } - - private def toSeq(value: Any): Seq[Any] = value match { - case null => Seq.empty - case values: java.lang.Iterable[_] => values.asScala.toSeq - case values: Iterable[_] => values.toSeq - case values: Array[_] => values.toSeq - case other => - throw new IllegalArgumentException( - s"Expected a collection in native Lance scan descriptor, got ${other.getClass.getName}") - } - - private def toStringMap(value: Any): Map[String, String] = value match { - case null => Map.empty - case values: java.util.Map[_, _] => - values.asScala.map { case (key, value) => key.toString -> value.toString }.toMap - case values: collection.Map[_, _] => - values.map { case (key, value) => key.toString -> value.toString }.toMap - case other => - throw new IllegalArgumentException( - s"Expected a map in native Lance scan descriptor, got ${other.getClass.getName}") - } - - private def toBoolean(value: Any): Boolean = value match { - case value: java.lang.Boolean => value.booleanValue() - case value: Boolean => value - case other => - throw new IllegalArgumentException( - s"Expected boolean in native Lance scan descriptor, got ${typeName(other)}") - } - - private def toLong(value: Any): Long = value match { - case value: java.lang.Number => value.longValue() - case value: String => value.toLong - case other => - throw new IllegalArgumentException( - s"Expected integer in native Lance scan descriptor, got ${typeName(other)}") - } - - private def toUInt32(value: Any, methodName: String): Int = { - val longValue = toLong(value) - if (longValue < 0 || longValue > 0xffffffffL) { - throw new IllegalArgumentException( - s"Native Lance scan descriptor method $methodName returned out-of-range uint32 " + - s"value $longValue") - } - longValue.toInt - } - - private def requireString(value: Any): String = value match { - case null => "" - case value: String => value - case other => other.toString - } - - private def nonEmptyString(value: Any, fallback: String): String = { - val stringValue = requireString(value) - if (stringValue.nonEmpty) stringValue else fallback - } - - private def typeName(value: Any): String = - Option(value).map(_.getClass.getName).getOrElse("null") - - private def scanKey(scanExec: BatchScanExec): String = - s"lance_${scanExec.id}_${scanExec.scan.hashCode()}" } diff --git a/spark/src/contrib-lance/scala/org/apache/spark/sql/comet/CometLanceNativeScanExec.scala b/spark/src/contrib-lance/scala/org/apache/spark/sql/comet/CometLanceNativeScanExec.scala index 5843e925364..4a24c21b0c8 100644 --- a/spark/src/contrib-lance/scala/org/apache/spark/sql/comet/CometLanceNativeScanExec.scala +++ b/spark/src/contrib-lance/scala/org/apache/spark/sql/comet/CometLanceNativeScanExec.scala @@ -44,7 +44,7 @@ case class CometLanceNativeScanExec( override val sourceKey: String, lanceDescriptor: LanceNativeScanDescriptor) extends CometLeafExec - with CometLanceNativeScanLike { + with CometScanWithPlanData { override val supportsColumnar: Boolean = true diff --git a/spark/src/contrib-lance/scala/org/apache/spark/sql/comet/LancePlanDataInjector.scala b/spark/src/contrib-lance/scala/org/apache/spark/sql/comet/LancePlanDataInjector.scala new file mode 100644 index 00000000000..ea002c1b4e4 --- /dev/null +++ b/spark/src/contrib-lance/scala/org/apache/spark/sql/comet/LancePlanDataInjector.scala @@ -0,0 +1,57 @@ +/* + * 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 org.apache.spark.sql.comet + +import org.apache.comet.serde.OperatorOuterClass +import org.apache.comet.serde.OperatorOuterClass.Operator +import org.apache.comet.serde.operator.CometLanceNativeScan + +class LancePlanDataInjector extends PlanDataInjector { + + override val opStructCase: Operator.OpStructCase = Operator.OpStructCase.CONTRIB_SCAN + + override def canInject(op: Operator): Boolean = + op.hasContribScan && + op.getContribScan.getTypeUrl == CometLanceNativeScan.LanceScanTypeUrl && + lanceScan(op).hasCommon && + !lanceScan(op).hasPartition + + override def getKey(op: Operator): Option[String] = + Some(lanceScan(op).getCommon.getScanId) + + override def inject( + op: Operator, + commonBytes: Array[Byte], + partitionBytes: Array[Byte]): Operator = { + val common = OperatorOuterClass.LanceScanCommon.parseFrom(commonBytes) + val partitionOnly = OperatorOuterClass.LanceScan.parseFrom(partitionBytes) + + val scanBuilder = lanceScan(op).toBuilder + scanBuilder.setCommon(common) + scanBuilder.setPartition(partitionOnly.getPartition) + + val contribBuilder = op.getContribScan.toBuilder + contribBuilder.setValue(scanBuilder.build().toByteString) + op.toBuilder.setContribScan(contribBuilder).build() + } + + private def lanceScan(op: Operator): OperatorOuterClass.LanceScan = + OperatorOuterClass.LanceScan.parseFrom(op.getContribScan.getValue) +} diff --git a/spark/src/main/scala/org/apache/comet/lance/LanceIntegration.scala b/spark/src/main/scala/org/apache/comet/lance/LanceIntegration.scala deleted file mode 100644 index 5424bb5f6d9..00000000000 --- a/spark/src/main/scala/org/apache/comet/lance/LanceIntegration.scala +++ /dev/null @@ -1,194 +0,0 @@ -/* - * 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 org.apache.comet.lance - -import java.lang.reflect.InvocationTargetException -import java.util.{Optional => JOptional} - -import scala.util.control.NonFatal - -import org.apache.spark.internal.Logging -import org.apache.spark.sql.execution.SparkPlan -import org.apache.spark.sql.execution.datasources.v2.BatchScanExec - -import org.apache.comet.CometConf -import org.apache.comet.CometSparkSessionExtensions.withFallbackReason - -/** - * Reflection-only bridge for optional Lance Spark integration. - * - * Default Comet builds must not depend on Lance classes. This object treats both Lance Spark and - * the Comet contrib-lance scaffold as optional runtime classes and falls back cleanly when either - * side is absent. - */ -object LanceIntegration extends Logging { - - private val LanceScanClassName = "org.lance.spark.read.LanceScan" - private val NativeScanPlanMethod = "nativeScanPlan" - private val ContribSupportModule = "org.apache.comet.lance.CometLanceSupport$" - - def isLanceScan(scan: Any): Boolean = { - scan != null && { - scan.getClass.getName == LanceScanClassName || - loadClass(LanceScanClassName).exists(_.isInstance(scan)) - } - } - - def nativeScanPlan(scan: Any): Option[Any] = - if (isLanceScan(scan)) { - invokeNativeScanPlan(scan) - } else { - None - } - - def tryCreateNativeScan(scanExec: BatchScanExec): Option[SparkPlan] = { - if (!CometConf.COMET_LANCE_NATIVE_ENABLED.get(scanExec.conf)) { - withFallbackReason( - scanExec, - s"Native Lance scan disabled because ${CometConf.COMET_LANCE_NATIVE_ENABLED.key} " + - "is not enabled") - return None - } - - if (!CometConf.COMET_EXEC_ENABLED.get(scanExec.conf)) { - withFallbackReason( - scanExec, - s"Native Lance scan disabled because ${CometConf.COMET_EXEC_ENABLED.key} is not enabled") - return None - } - - val nativePlan = nativeScanPlan(scanExec.scan) match { - case Some(plan) => plan - case None => - withFallbackReason( - scanExec, - s"Native Lance scan disabled because $LanceScanClassName.$NativeScanPlanMethod() " + - "is not available") - return None - } - - val support = loadContribSupport match { - case Some(module) => module - case None => - withFallbackReason( - scanExec, - "Native Lance scan disabled because the contrib-lance build profile is not present") - return None - } - - try { - val method = - support.getClass.getMethod("tryTransform", classOf[BatchScanExec], classOf[Object]) - method.invoke(support, scanExec, nativePlan.asInstanceOf[AnyRef]) match { - case plan: Option[_] => plan.asInstanceOf[Option[SparkPlan]] - case other => - logWarning( - "Native Lance scan disabled because contrib-lance returned unexpected " + - s"result: ${Option(other).map(_.getClass.getName).getOrElse("null")}") - None - } - } catch { - case e: InvocationTargetException => - val cause = Option(e.getCause).getOrElse(e) - logWarning( - "Native Lance scan disabled because contrib-lance threw during reflection: " + - s"${cause.getClass.getName}: ${cause.getMessage}", - cause) - None - case NonFatal(e) => - logWarning(s"Native Lance scan disabled by contrib-lance reflection failure: $e") - None - } - } - - private[comet] def invokeNativeScanPlan(scan: Any): Option[Any] = { - try { - findNoArgMethod(scan.getClass, NativeScanPlanMethod) - .flatMap { method => - optionalResult(method.invoke(scan)) - } - } catch { - case e: InvocationTargetException => - logWarning( - s"Native Lance scan disabled because $NativeScanPlanMethod() threw: " + - s"${Option(e.getCause).map(_.getMessage).getOrElse(e.getMessage)}") - None - case NonFatal(e) => - logWarning(s"Native Lance scan disabled by reflection failure: $e") - None - } - } - - private def optionalResult(value: Any): Option[Any] = value match { - case null => None - case option: Option[_] => option - case option: JOptional[_] if option.isPresent => Some(option.get) - case _: JOptional[_] => None - case other => Some(other) - } - - private def findNoArgMethod( - clazz: Class[_], - methodName: String): Option[java.lang.reflect.Method] = { - var current = clazz - while (current != null) { - try { - val method = current.getDeclaredMethod(methodName) - method.setAccessible(true) - return Some(method) - } catch { - case _: NoSuchMethodException => - current = current.getSuperclass - case NonFatal(_) => - return None - } - } - None - } - - private def loadContribSupport: Option[AnyRef] = - loadClass(ContribSupportModule).flatMap { clazz => - try { - Some(clazz.getField("MODULE$").get(null).asInstanceOf[AnyRef]) - } catch { - case NonFatal(_) => None - } - } - - private def loadClass(className: String): Option[Class[_]] = { - try { - val classLoader = Thread.currentThread().getContextClassLoader - // scalastyle:off classforname - val clazz = - if (classLoader != null) { - Class.forName(className, false, classLoader) - } else { - Class.forName(className) - } - // scalastyle:on classforname - Some(clazz) - } catch { - case _: ClassNotFoundException | _: NoClassDefFoundError => None - case NonFatal(e) => - logDebug(s"Unable to load optional class $className", e) - None - } - } -} diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometLanceNativeScanLike.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometLanceNativeScanLike.scala deleted file mode 100644 index 72b5326b7b9..00000000000 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometLanceNativeScanLike.scala +++ /dev/null @@ -1,28 +0,0 @@ -/* - * 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 org.apache.spark.sql.comet - -private[comet] trait CometLanceNativeScanLike extends CometLeafExec { - def sourceKey: String - - def commonData: Array[Byte] - - def perPartitionData: Array[Array[Byte]] -} diff --git a/spark/src/test/scala/org/apache/comet/rules/CometScanRuleSuite.scala b/spark/src/test/scala/org/apache/comet/rules/CometScanRuleSuite.scala index 15a9de3dde3..2ac0ac0b0fb 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometScanRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometScanRuleSuite.scala @@ -202,8 +202,7 @@ class CometScanRuleSuite extends CometTestBase { val projectedSchema = StructType(Seq(StructField("id", DataTypes.IntegerType, false))) val descriptor = new FakeLanceNativeScanPlan(requiredSchema, projectedSchema) - val (common, partitions) = - serializeFakeLanceDescriptor(serde, descriptor, "fallback-scan", requiredSchema) + val (common, partitions) = serializeFakeLanceDescriptor(serde, descriptor) assert(common.getScanId == "scan-123") assert(common.getDatasetUri == "s3://bucket/table.lance") @@ -238,21 +237,17 @@ class CometScanRuleSuite extends CometTestBase { .asInstanceOf[AnyRef] }.toOption - private def serializeFakeLanceDescriptor( - serde: AnyRef, - descriptor: AnyRef, - fallbackScanId: String, - fallbackRequiredSchema: StructType) + private def serializeFakeLanceDescriptor(serde: AnyRef, descriptor: AnyRef) : (OperatorOuterClass.LanceScanCommon, Array[OperatorOuterClass.LanceScan]) = { val method = serde.getClass.getMethods .find(method => - method.getName == "serializeNativePlan" && method.getParameterTypes.length == 3) + method.getName == "serializeNativePlan" && method.getParameterTypes.length == 1) .getOrElse { throw new AssertionError("CometLanceNativeScan.serializeNativePlan was not found") } val serialized = method - .invoke(serde, descriptor, fallbackScanId, fallbackRequiredSchema) + .invoke(serde, descriptor) .asInstanceOf[Product] val commonBytes = serialized.productElement(0).asInstanceOf[Array[Byte]] val partitionBytes = serialized.productElement(1).asInstanceOf[Array[Array[Byte]]] @@ -262,9 +257,7 @@ class CometScanRuleSuite extends CometTestBase { partitionBytes.map(OperatorOuterClass.LanceScan.parseFrom)) } - private class FakeLanceNativeScanPlan( - requiredSchema: StructType, - projectedSchema: StructType) { + private class FakeLanceNativeScanPlan(requiredSchema: StructType, projectedSchema: StructType) { private val storageOptions = new LinkedHashMap[String, String]() storageOptions.put("region", "us-west-2") storageOptions.put("endpoint", "http://127.0.0.1:9000")