gengliangwang commented on code in PR #59520: URL: https://github.com/apache/spark/pull/59520#discussion_r4497421976
########## sql/native/rust/src/bridge.rs: ########## @@ -0,0 +1,876 @@ +// 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. + +//! Driver-only handles of sources, scans, streams, and writes, and executor readers restored +//! from serialized values. + +use std::any::Any; +use std::collections::HashMap; +use std::ffi::CString; +use std::panic::{catch_unwind, AssertUnwindSafe}; +use std::sync::atomic::{AtomicI64, Ordering}; +use std::sync::{Arc, Mutex, MutexGuard}; + +use arrow_array::ffi_stream::FFI_ArrowArrayStream; +use arrow_schema::SchemaRef; +use serde::{de::DeserializeOwned, Deserialize, Serialize}; + +use crate::write::{Batch, DriverWrite, DriverWriter}; +use crate::{ + abi, BatchWrite, Capabilities, CreateInfo, DataSource, LatestOffsetInfo, LimitSupport, + MicroBatchStream, MicroBatchStreamInfo, PartitionReaderFactory, ReadLimit, ReaderContext, + Result, ScanBuilder, ScanInfo, ScanPlan, StreamingWrite, StreamingWriteInfo, TaskContext, + WriteInfo, WriteMode, +}; + +#[path = "stream.rs"] +pub(crate) mod stream; + +const STATE_VERSION: u32 = 3; + +type HandleMap<T> = Mutex<HashMap<i64, Arc<Mutex<T>>>>; + +/// One registry per loaded plugin. Provider code runs under a per-handle lock, +/// so independent sources, scans, and streams do not serialize behind a global lock. +pub struct Bridge<S: DataSource> { + next_handle: AtomicI64, + sources: HandleMap<SourceState<S>>, + // The builder of each scan, until Spark plans the scan. + scans: HandleMap<Option<S::ScanBuilder>>, + streams: HandleMap<S::MicroBatchStream>, + writers: HandleMap<Box<dyn DriverWrite>>, +} + +impl<S: DataSource> Default for Bridge<S> { + fn default() -> Self { + Self { + next_handle: AtomicI64::new(1), + sources: Mutex::new(HashMap::new()), + scans: Mutex::new(HashMap::new()), + streams: Mutex::new(HashMap::new()), + writers: Mutex::new(HashMap::new()), + } + } +} + +/// The partitions of a scan or a micro-batch, and the state of its reader: its factory. +pub struct EncodedPlan { + pub state: Vec<u8>, + pub partitions: Vec<Vec<u8>>, +} + +impl<S: DataSource> Bridge<S> { + // Whether the source implements each operation, whose native callbacks the SDK installs. + + pub const READ: bool = <S::ScanBuilder as ScanBuilder>::SUPPORTED; + + pub const MICRO_BATCH_READ: bool = <S::MicroBatchStream as MicroBatchStream>::SUPPORTED; + + pub const WRITE: bool = <S::BatchWrite as BatchWrite>::SUPPORTED; + + pub const STREAMING_WRITE: bool = <S::StreamingWrite as StreamingWrite>::SUPPORTED; + + const CAPABILITIES: Capabilities = Capabilities::implemented(&[ + (Capabilities::BATCH_READ, Self::READ), + (Capabilities::MICRO_BATCH_READ, Self::MICRO_BATCH_READ), + (Capabilities::BATCH_WRITE, Self::WRITE), + (Capabilities::STREAMING_WRITE, Self::STREAMING_WRITE), + ( + Capabilities::BATCH_TRUNCATE, + <S::BatchWrite as BatchWrite>::SUPPORTS_TRUNCATE, + ), + ( + Capabilities::STREAMING_TRUNCATE, + <S::StreamingWrite as StreamingWrite>::SUPPORTS_TRUNCATE, + ), + ]); + + fn insert<T>(&self, handles: &HandleMap<T>, value: T) -> Result<i64> { + let handle = self + .next_handle + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |id| id.checked_add(1)) + .map_err(|_| "native handle IDs exhausted")?; + lock(handles).insert(handle, Arc::new(Mutex::new(value))); + Ok(handle) + } + + pub fn create(&self, info: CreateInfo) -> Result<(i64, u64)> { + self.register_source(S::new(info)?) + } + + fn register_source(&self, source: S) -> Result<(i64, u64)> { + let capabilities = Self::effective(source.capabilities()); + let handle = self.insert( + &self.sources, + SourceState { + source, + capabilities, + }, + )?; + // Every stream supports admission control: its default read limit is all the + // available data unless it overrides MicroBatchStream::default_read_limit. + let admission_control = if capabilities.contains(Capabilities::MICRO_BATCH_READ) { + abi::ADMISSION_CONTROL + } else { + 0 + }; + Ok((handle, capabilities.bits() | admission_control)) + } + + /// Keeps the implemented operations of a declaration, so that a source can narrow its + /// operations but never add one, and each truncation only with its write. + fn effective(declared: Capabilities) -> Capabilities { + let mut capabilities = declared.intersection(Self::CAPABILITIES); + for (write, truncate) in [ + (Capabilities::BATCH_WRITE, Capabilities::BATCH_TRUNCATE), + ( + Capabilities::STREAMING_WRITE, + Capabilities::STREAMING_TRUNCATE, + ), + ] { + if !capabilities.contains(write) { + capabilities = capabilities.without(truncate); + } + } + capabilities + } + + pub fn schema(&self, handle: i64) -> Result<SchemaRef> { + let source = get(&self.sources, handle)?; + let schema = lock(&source).source.schema()?; + stream::validate_schema(&schema)?; + Ok(schema) + } + + pub fn create_reader(&self, handle: i64, schema: SchemaRef) -> Result<i64> { + stream::validate_schema(&schema)?; + let source = get(&self.sources, handle)?; + let builder = { + let source = lock(&source); + require(source.capabilities, Capabilities::BATCH_READ, "batch reads")?; + source.source.new_scan(ScanInfo::new(schema))? + }; + self.insert(&self.scans, Some(builder)) + } + + pub fn push_limit(&self, handle: i64, limit: usize) -> Result<bool> { + let scan = get(&self.scans, handle)?; + let mut scan = lock(&scan); + Ok(builder(&mut scan)?.push_limit(limit)? == LimitSupport::Global) + } + + pub fn prune_columns(&self, handle: i64, columns: &[String]) -> Result<()> { + let scan = get(&self.scans, handle)?; + let mut scan = lock(&scan); + builder(&mut scan)?.prune_columns(columns) + } + + /// Plans a scan, which consumes its builder. + pub fn plan(&self, handle: i64) -> Result<EncodedPlan> { + let scan = get(&self.scans, handle)?; + let builder = lock(&scan) + .take() + .ok_or("the scan has already been planned")?; + encode_plan(builder.plan()?) + } + + pub fn close_reader(&self, handle: i64) -> Result<()> { + remove(&self.scans, handle); + Ok(()) + } + + pub fn create_stream_reader(&self, handle: i64, info: MicroBatchStreamInfo) -> Result<i64> { + stream::validate_schema(&info.schema)?; + let source = get(&self.sources, handle)?; + let stream = { + let source = lock(&source); + require( + source.capabilities, + Capabilities::MICRO_BATCH_READ, + "micro-batch streaming reads", + )?; + source.source.new_micro_batch_stream(info)? + }; + self.insert(&self.streams, stream) + } + + pub fn initial_offset(&self, handle: i64) -> Result<String> { + let stream = get(&self.streams, handle)?; + let offset = lock(&stream).initial_offset()?; + encode_offset(&offset) + } + + pub fn default_read_limit(&self, handle: i64) -> Result<ReadLimit> { + let stream = get(&self.streams, handle)?; + let limit = lock(&stream).default_read_limit(); + limit + } + + /// Decodes the offsets of a request, and encodes the latest offset. + pub fn latest_offset(&self, handle: i64, info: LatestOffsetInfo<String>) -> Result<String> { + let stream = get(&self.streams, handle)?; + let info = LatestOffsetInfo { + start: decode_offset(&info.start)?, + limit: info.limit, + run_end: info.run_end.as_deref().map(decode_offset).transpose()?, + }; + let offset = lock(&stream).latest_offset(info)?; + encode_offset(&offset) + } Review Comment: Done in e933ef2cf382: The `MicroBatchStream` offset docs and the README now say that a changed `Offset` type can re-encode the same position differently, which makes Spark run one micro-batch without rows. ########## sql/native/rust/src/stream.rs: ########## @@ -0,0 +1,830 @@ +// 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::ffi::{c_char, c_int, c_void, CString}; +use std::sync::Arc; + +use arrow_array::ffi::FFI_ArrowArray; +use arrow_array::ffi_stream::FFI_ArrowArrayStream; +use arrow_array::{Array, RecordBatch, StructArray}; +use arrow_data::ArrayData; +use arrow_schema::{ffi::FFI_ArrowSchema, DataType, Field, Fields, Schema, SchemaRef, TimeUnit}; + +use super::{c_message, drop_provider, provider_call}; +use crate::{PartitionReader, Result}; + +struct StreamState<'task> { + reader: Option<PartitionReader<'task>>, + // The schema to read, whose columns, Spark types, and nullability each batch must have. + expected: SchemaRef, + // The schema of the stream: that of the first batch, which every batch must have, or + // the schema to read of a reader without batches. Unknown until either is read. + schema: Option<SchemaRef>, + // The first batch, when get_schema read it to learn the schema. + pending: Option<ArrayData>, + last_error: Option<CString>, + finished: bool, +} + +impl StreamState<'_> { + /// Reads and checks the next batch. The first batch, or the end of the reader, + /// determines the schema of the stream. + fn read_next(&mut self) -> Result<Option<ArrayData>> { + let Some(batch) = self.reader.as_mut().expect("reader is present").next() else { + self.finished = true; + self.schema.get_or_insert_with(|| self.expected.clone()); + return Ok(None); + }; + let batch = project(batch?, &self.expected)?; + match &self.schema { + Some(schema) if batch.schema() != *schema => { + return Err("reader returned a batch with a different schema".into()); Review Comment: Done in e933ef2cf382: The schema-change error now names the first column whose field differs, or says the schema metadata differs. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
