diff --git a/Cargo.lock b/Cargo.lock index 0b9abbd87..bac3ba8c1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2386,6 +2386,7 @@ dependencies = [ "insta", "krishiv-common", "krishiv-plan", + "krishiv-sql", "krishiv-types", "prost", "prost-reflect", diff --git a/changelog.d/774.added.md b/changelog.d/774.added.md new file mode 100644 index 000000000..0534f4158 --- /dev/null +++ b/changelog.d/774.added.md @@ -0,0 +1 @@ +Spark Connect translates local relations, declared file schemas, qualified subquery aliases, nested literals, expression strings and `spark.sql()` (sql_command), as part of group 1 (PR-12 part 2b, #774). diff --git a/crates/krishiv-connect-server/CLAUDE.md b/crates/krishiv-connect-server/CLAUDE.md index 83efc0420..b637f7862 100644 --- a/crates/krishiv-connect-server/CLAUDE.md +++ b/crates/krishiv-connect-server/CLAUDE.md @@ -50,6 +50,9 @@ Paths are relative to this crate. operation's `Replay` under the ahead budget (`krishiv.connect.execute.aheadBytes`) and memory consumer `connect_responses` (the Arrow payload is `Bytes`, shared with the copy sent), the response order (schema, batches, metrics, `ResultComplete` when reattachable); commands are 56K00 until T012; + `request_plan` (shared with `AnalyzePlan`) translates twice when the plan holds local relations: the first pass records + them, `Session::local_relation` builds them (cached ones are `NOT_FOUND_CACHED_LOCAL_RELATION` until the artifact + store lands), the second pass resolves them (`krishiv-connect` CLAUDE.md); `src/execute/arrow.rs` Spark's Arrow shapes, batch limits, IPC streams, codecs and chunks; `src/execute/schema.rs` the `schema` response (`DataTypeProtoConverter.toConnectProtoType`) - `src/codec.rs` `WireCodec`: requests decode only through `wire::decode_request` (bounded size, count, depth, fields) diff --git a/crates/krishiv-connect-server/src/analyze.rs b/crates/krishiv-connect-server/src/analyze.rs index 57babb56a..9bf5c8ce6 100644 --- a/crates/krishiv-connect-server/src/analyze.rs +++ b/crates/krishiv-connect-server/src/analyze.rs @@ -7,7 +7,7 @@ //! - `schema`: the analyzed schema, the same `DataType` `ExecutePlan` sends; //! - `explain`: `DataFrame::explain` in the requested mode; //! - `tree_string`: Spark's `StructType.treeString` (`analyze::types`); -//! - `is_local`: whether the translated root is an inline table (Spark's +//! - `is_local`: whether the translated root is a local relation (Spark's //! `LocalRelation`); `is_streaming`: always false (no streaming sources); //! - `input_files`: the paths of the file scans, as written; //! - `spark_version`: the session profile's release (`4.2.0` for 4.2); @@ -68,13 +68,13 @@ pub async fn analyze_plan( /// recurse on the async worker's stack (translation runs a deep plan on a /// big stack of its own). async fn frame( - session: &krishiv_session::Session, + entry: &SessionEntry, plan: Option<&pb::Plan>, limits: &ExecuteLimits, ) -> Result { let empty = pb::Plan::default(); - let logical = request_plan(session, plan.unwrap_or(&empty), limits).await?; - Ok(session.from_plan(&logical)) + let logical = request_plan(entry, plan.unwrap_or(&empty), limits).await?; + Ok(entry.session().from_plan(&logical)) } async fn analyze( @@ -91,7 +91,7 @@ async fn analyze( .into()); } Some(Analyze::Schema(s)) => { - let df = frame(session, s.plan.as_ref(), limits).await?; + let df = frame(entry, s.plan.as_ref(), limits).await?; res::Result::Schema(res::Schema { schema: Some(schema::proto_schema(&*df.schema().await?)?), }) @@ -99,7 +99,7 @@ async fn analyze( Some(Analyze::Explain(e)) => { // Spark builds the Dataset (translation, analysis) before it // reads the mode. - let df = frame(session, e.plan.as_ref(), limits).await?; + let df = frame(entry, e.plan.as_ref(), limits).await?; df.analyzed().await?; let mode = explain_mode(e.explain_mode)?; res::Result::Explain(res::Explain { @@ -107,29 +107,31 @@ async fn analyze( }) } Some(Analyze::TreeString(t)) => { - let df = frame(session, t.plan.as_ref(), limits).await?; + let df = frame(entry, t.plan.as_ref(), limits).await?; let spark = schema::spark_schema(&*df.schema().await?)?; res::Result::TreeString(res::TreeString { tree_string: types::tree_string(&spark, t.level), }) } Some(Analyze::IsLocal(l)) => { - // Spark: the analyzed plan is a `LocalRelation`. Krishiv's - // analyzer lowers inline rows to `OneRowRelation` (as for - // `SELECT 1`), so the question is asked of the translated plan - // (an inline table, as `LocalRelation` will translate in T012); - // SQL `VALUES` is SQL text until analysis and answers false. - // The plan is still analyzed so that an invalid one fails as in - // Spark. + // Spark: the analyzed plan is a `LocalRelation`. A `local_relation` translates to + // a `Read` of a local relation (its user schema already applied eagerly, so the + // analyzed plan is that bare relation either way); Krishiv's analyzer lowers inline + // rows to `OneRowRelation` (as for `SELECT 1`), so the question is asked of the + // translated plan; SQL `VALUES` is SQL text until analysis and answers false. The + // plan is still analyzed so that an invalid one fails as in Spark. let empty = pb::Plan::default(); - let logical = request_plan(session, l.plan.as_ref().unwrap_or(&empty), limits).await?; - let is_local = matches!(logical.root().node, PlanNode::InlineTable { .. }); + let logical = request_plan(entry, l.plan.as_ref().unwrap_or(&empty), limits).await?; + let is_local = matches!( + logical.root().node, + PlanNode::InlineTable { .. } | PlanNode::Read(ReadSource::Local(_)) + ); session.from_plan(&logical).analyzed().await?; res::Result::IsLocal(res::IsLocal { is_local }) } Some(Analyze::IsStreaming(s)) => { // Analyzed so that an invalid plan fails as in Spark. - frame(session, s.plan.as_ref(), limits) + frame(entry, s.plan.as_ref(), limits) .await? .analyzed() .await?; @@ -138,7 +140,7 @@ async fn analyze( }) } Some(Analyze::InputFiles(f)) => { - let analyzed = frame(session, f.plan.as_ref(), limits) + let analyzed = frame(entry, f.plan.as_ref(), limits) .await? .analyzed() .await?; @@ -157,15 +159,14 @@ async fn analyze( )?)?), }), Some(Analyze::SameSemantics(s)) => { - let target = - semantic_hash(frame(session, s.target_plan.as_ref(), limits).await?).await?; - let other = semantic_hash(frame(session, s.other_plan.as_ref(), limits).await?).await?; + let target = semantic_hash(frame(entry, s.target_plan.as_ref(), limits).await?).await?; + let other = semantic_hash(frame(entry, s.other_plan.as_ref(), limits).await?).await?; res::Result::SameSemantics(res::SameSemantics { result: target == other, }) } Some(Analyze::SemanticHash(s)) => { - let hash = semantic_hash(frame(session, s.plan.as_ref(), limits).await?).await?; + let hash = semantic_hash(frame(entry, s.plan.as_ref(), limits).await?).await?; res::Result::SemanticHash(res::SemanticHash { result: fold_hash(hash), }) diff --git a/crates/krishiv-connect-server/src/analyze/types.rs b/crates/krishiv-connect-server/src/analyze/types.rs index 2ad24f1ea..998fad0d0 100644 --- a/crates/krishiv-connect-server/src/analyze/types.rs +++ b/crates/krishiv-connect-server/src/analyze/types.rs @@ -183,28 +183,7 @@ pub fn json_to_ddl(json: &str) -> Result { } } -/// The planner's `parseDatatypeString` (`SparkConnectPlanner.scala:1511-1528`): -/// a table schema (`a INT, b STRING`), else a data type (`int`), else the -/// text as the fields of a struct; the first failure is the error. -/// -/// Krishiv parses the table schema as the fields of a `STRUCT<...>`, so -/// clauses only a table schema takes (`DEFAULT`, `GENERATED`) are refused, -/// and the error is the data type's (Spark reports the table schema's): a -/// recorded difference. -/// -/// # Errors -/// The parse error of the text as a data type. -pub fn parse_datatype_string(text: &str) -> Result { - if let Ok(t) = krishiv_types::parse_data_type(&format!("STRUCT<{text}>")) { - return Ok(t); - } - match krishiv_types::parse_data_type(text) { - Ok(t) => Ok(t), - Err(e) => { - krishiv_types::parse_data_type(&format!("STRUCT<{}>", text.trim())).map_err(|_| e) - } - } -} +pub use krishiv_connect::translate::parse_datatype_string; #[cfg(test)] mod tests; diff --git a/crates/krishiv-connect-server/src/execute.rs b/crates/krishiv-connect-server/src/execute.rs index 40e7a9b74..238b14c2c 100644 --- a/crates/krishiv-connect-server/src/execute.rs +++ b/crates/krishiv-connect-server/src/execute.rs @@ -27,8 +27,9 @@ //! //! Errors end the stream with the gRPC status of [`crate::status`] after the //! responses before them. An interrupted operation ends with -//! `OPERATION_CANCELED`. Commands answer `CONNECT_INVALID_PLAN` (56K00) -//! until P3-T012. +//! `OPERATION_CANCELED`. A `sql_command` (`spark.sql()`) answers one +//! `sql_command_result` (`execute::sql_command`); the other commands answer +//! `CONNECT_INVALID_PLAN` (56K00) until P3-T012. //! //! The core is protocol-neutral where it is cheap: the plan is decoded into //! a `krishiv_plan::Plan`, run as a `DataFrame` and encoded as Arrow IPC; @@ -36,6 +37,7 @@ pub mod arrow; pub mod schema; +mod sql_command; use std::io::Read; use std::sync::Arc; @@ -44,8 +46,8 @@ use arrow_schema::SchemaRef; use futures::StreamExt; use krishiv_common::error::{ ConnectInvalidPlanInvalidEnum, ConnectInvalidPlanInvalidOneOfFieldNotSet, - ConnectInvalidPlanPlanSizeLargerThanMax, InvalidCursorNotReattachable, KrishivError, - OperationCanceled, + ConnectInvalidPlanNotFoundCachedLocalRelation, ConnectInvalidPlanPlanSizeLargerThanMax, + InvalidCursorNotReattachable, KrishivError, OperationCanceled, }; use krishiv_connect::proto::v4_3 as pb; use krishiv_connect::proto::v4_3::execute_plan_request::request_option::RequestOption; @@ -54,7 +56,9 @@ use krishiv_connect::proto::v4_3::execute_plan_response::{ }; use krishiv_connect::proto::v4_3::plan::OpType; use krishiv_connect::proto::v4_3::plan::compressed_operation::OpType as CompressedOpType; -use krishiv_connect::translate::{TranslateConf, command_not_supported, translate_relation}; +use krishiv_connect::translate::{ + LocalRelations, LocalRequest, TranslateConf, command_not_supported, translate_relation, +}; use krishiv_connect::wire; use krishiv_memory::ConsumerHandle; use krishiv_session::{DataFrame, QueryOptions, RecordBatchStream}; @@ -258,17 +262,31 @@ pub async fn execute_plan( let empty = pb::Plan::default(); let plan = request.plan.as_ref().unwrap_or(&empty); let prepared = until_cancelled(&op, async { + // `spark.sql()`: a command runs now, a query is analyzed; neither returns rows. + if let Some(command) = sql_command::of(plan, limits).await? { + let answer = + sql_command::run(session, &command, id.as_str(), limits.max_batch_size).await?; + return Ok(Prepared::Command(answer)); + } // Decompression (up to `maxPlanSize`) runs off the async workers. - let logical = request_plan(session, plan, limits).await?; + let logical = request_plan(entry, plan, limits).await?; let df = session.from_plan_with(&logical, QueryOptions::new().operation_id(id.as_str()))?; let schema = df.schema().await?; - Ok((df, schema)) + Ok(Prepared::Rows(df, schema)) }) .await; // A preparation error settles like a query error: an interrupt that // reported this operation first turns it into `OPERATION_CANCELED`. let (df, schema) = match prepared { - Ok(ready) => ready, + Ok(Prepared::Rows(df, schema)) => (df, schema), + Ok(Prepared::Command(answer)) => { + let sender = new_sender(session, request, entry, &op, (limits, budget, &opts)); + let stream = + sender::attach(Arc::clone(&op), operations, 0, SenderCaps::of(&op, limits)); + registered.keep(); + tokio::spawn(sql_command::encode(sender, answer)); + return Ok(stream); + } Err(e) => { return Err(if op.finish() { e @@ -280,7 +298,32 @@ pub async fn execute_plan( let schema_response = schema::proto_schema(&schema)?; let settings = encode_settings(&conf, limits, opts.chunking); - let sender = Sender { + let sender = new_sender(session, request, entry, &op, (limits, budget, &opts)); + // Attached before the encoder starts: the operation is never detached + // before its first stream. + let stream = sender::attach(Arc::clone(&op), operations, 0, SenderCaps::of(&op, limits)); + registered.keep(); + tokio::spawn(encode(sender, df, schema, schema_response, settings)); + Ok(stream) +} + +/// What a request's plan prepared. +enum Prepared { + /// A relation: its frame and schema, streamed as batches. + Rows(DataFrame, SchemaRef), + /// A `sql_command`: the relation it answers. + Command(sql_command::Answer), +} + +/// The producing side of `op`. +fn new_sender( + session: &krishiv_session::Session, + request: &wire::Decoded, + entry: &SessionEntry, + op: &Arc, + (limits, budget, opts): (&ExecuteLimits, &ResponseBudget, &Options), +) -> Sender { + Sender { ahead: Arc::new(Semaphore::new( limits.ahead_bytes.clamp(1, Semaphore::MAX_PERMITS), )), @@ -292,16 +335,10 @@ pub async fn execute_plan( 0 }), consumer: ConsumerHandle::register(session.runtime().memory(), "connect_responses", None), - op: Arc::clone(&op), + op: Arc::clone(op), session_id: request.session_id.clone(), server_side_session_id: entry.server_side_id().to_owned(), - }; - // Attached before the encoder starts: the operation is never detached - // before its first stream. - let stream = sender::attach(Arc::clone(&op), operations, 0, SenderCaps::of(&op, limits)); - registered.keep(); - tokio::spawn(encode(sender, df, schema, schema_response, settings)); - Ok(stream) + } } /// The engine plan of a request's `Plan`: a compressed one decompressed @@ -313,15 +350,18 @@ pub async fn execute_plan( /// # Errors /// The decompression, decoding and translation errors. pub(crate) async fn request_plan( - session: &krishiv_session::Session, + entry: &SessionEntry, plan: &pb::Plan, limits: &ExecuteLimits, ) -> Result { + let session = entry.session(); let translate = TranslateConf { default_data_source: session .conf() .get("spark.sql.sources.default") .unwrap_or_else(|| TranslateConf::default().default_data_source), + parser: session.parser_options(), + local_relations: LocalRelations::default(), }; let decompressed = match &plan.op_type { Some(OpType::CompressedOperation(c)) => { @@ -340,7 +380,60 @@ pub(crate) async fn request_plan( } _ => None, }; - translate_plan(plan, decompressed.as_deref(), &translate) + // The local relations (`createDataFrame`) are built by an async pre-pass, as the translator is + // synchronous: a first translation records what the plan needs, the relations are built, and a + // second translation uses them. A plan without any is translated once. + let collecting = LocalRelations::collecting(); + let first = TranslateConf { + local_relations: collecting.clone(), + ..translate.clone() + }; + let logical = translate_plan(plan, decompressed.as_deref(), &first)?; + // Drained, not copied: the recorded inline data is the one copy, and each request is consumed + // (and freed) by its build below. + let requests = collecting.take_requests(); + drop(first); + drop(collecting); + if requests.is_empty() { + return Ok(logical); + } + drop(logical); + let mut built = Vec::with_capacity(requests.len()); + for request in requests { + built.push(build_local_relation(entry, request).await?); + } + let second = TranslateConf { + local_relations: LocalRelations::resolved(built), + ..translate + }; + translate_plan(plan, decompressed.as_deref(), &second) +} + +/// The relation of one request of the translator's first pass: an inline one through +/// `Session::local_relation` (the IPC data decoded, a user schema applied eagerly, the batches +/// held once in the session's relation store); a cached one through the artifact store. +/// +/// The artifact store is not on this branch (T011/T015b), so no hash is ever found: +/// `CONNECT_INVALID_PLAN.NOT_FOUND_CACHED_LOCAL_RELATION`, as Spark's block read of an unknown +/// `cache/` (golden-pending: Spark's `CANNOT_FIND_CACHED_LOCAL_RELATION` for a hash the +/// artifact manager never registered). +async fn build_local_relation( + entry: &SessionEntry, + request: LocalRequest, +) -> Result { + match request { + LocalRequest::Inline { data, schema } => { + entry + .session() + .local_relation(data.as_deref(), schema) + .await + } + LocalRequest::Cached { hash } => Err(ConnectInvalidPlanNotFoundCachedLocalRelation { + hash, + session_u_u_i_d: entry.server_side_id().to_owned(), + } + .into()), + } } /// The plan as the engine's IR: a ROOT relation translated, a compressed @@ -359,22 +452,10 @@ fn translate_plan( Ok(CompressedOpType::Relation) => { translate_relation(&*wire::decode_relation(bytes)?, conf) } - Ok(CompressedOpType::Command) => { - // Decoded through the bounded decoder as the `command` - // field (2) of a `Plan`, which `wire` can start from. - let mut wrapped = Vec::with_capacity(bytes.len() + 11); - prost::encoding::encode_key( - 2, - prost::encoding::WireType::LengthDelimited, - &mut wrapped, - ); - prost::encoding::encode_varint(bytes.len() as u64, &mut wrapped); - wrapped.extend_from_slice(bytes); - match &wire::decode_plan(&wrapped)?.op_type { - Some(OpType::Command(command)) => Err(command_not_supported(command)), - _ => Err(command_not_supported(&pb::Command::default())), - } - } + Ok(CompressedOpType::Command) => match &decode_command(bytes)?.op_type { + Some(OpType::Command(command)) => Err(command_not_supported(command)), + _ => Err(command_not_supported(&pb::Command::default())), + }, _ => { // Spark's `PlanCompressionUtils` re-throws what parsing // the operation raised as `CANNOT_PARSE` (v4.3.0). @@ -398,6 +479,16 @@ fn translate_plan( } } +/// A decompressed `Command`, decoded through the bounded decoder as the `command` field (2) of a +/// `Plan`, which `wire` can start from. +fn decode_command(bytes: &[u8]) -> Result, KrishivError> { + let mut wrapped = Vec::with_capacity(bytes.len() + 11); + prost::encoding::encode_key(2, prost::encoding::WireType::LengthDelimited, &mut wrapped); + prost::encoding::encode_varint(bytes.len() as u64, &mut wrapped); + wrapped.extend_from_slice(bytes); + wire::decode_plan(&wrapped) +} + /// Spark's `PlanCompressionUtils.decompressBytesWithZstd`: the frame's /// declared size is checked first, then the output is read through a bound, /// so a small input that expands past `max` fails before `max` bytes are diff --git a/crates/krishiv-connect-server/src/execute/arrow.rs b/crates/krishiv-connect-server/src/execute/arrow.rs index 998e7a04b..a7314c342 100644 --- a/crates/krishiv-connect-server/src/execute/arrow.rs +++ b/crates/krishiv-connect-server/src/execute/arrow.rs @@ -339,14 +339,32 @@ fn element_sizes(offsets: impl Iterator, values: &dyn Array) -> Ve pub fn ipc_stream( batch: &RecordBatch, compression: Option, +) -> Result, KrishivError> { + ipc_stream_of(std::slice::from_ref(batch), compression) +} + +/// `batches` (at least one: its schema is the stream's) as one Arrow IPC stream. +/// +/// # Errors +/// `INTERNAL_ERROR` when Arrow cannot write it, or no batch is given. +pub fn ipc_stream_of( + batches: &[RecordBatch], + compression: Option, ) -> Result, KrishivError> { let arrow = |e: arrow_schema::ArrowError| internal(format!("writing a result batch: {e}")); let options = IpcWriteOptions::default() .try_with_compression(compression) .map_err(arrow)?; + let Some(first) = batches.first() else { + return Err(internal( + "writing a result stream without a batch".to_owned(), + )); + }; let mut writer = - StreamWriter::try_new_with_options(Vec::new(), &batch.schema(), options).map_err(arrow)?; - writer.write(batch).map_err(arrow)?; + StreamWriter::try_new_with_options(Vec::new(), &first.schema(), options).map_err(arrow)?; + for batch in batches { + writer.write(batch).map_err(arrow)?; + } writer.finish().map_err(arrow)?; writer.into_inner().map_err(arrow) } diff --git a/crates/krishiv-connect-server/src/execute/sql_command.rs b/crates/krishiv-connect-server/src/execute/sql_command.rs new file mode 100644 index 000000000..3d80b820d --- /dev/null +++ b/crates/krishiv-connect-server/src/execute/sql_command.rs @@ -0,0 +1,202 @@ +// SPDX-License-Identifier: Apache-2.0 +//! `spark.sql()`: the `SqlCommand` of `Plan.command` (design P3-connect-mapping, command 5). +//! +//! The statement runs through [`Session::sql_command`]: a command (`SET`, `CREATE VIEW`, `SHOW`, +//! `INSERT`, ...) runs at once and answers a `LocalRelation` of its result rows; a query is only +//! analyzed (its errors surface here) and answers the `SQL` relation the client sends again as a +//! root relation. The stream is one `sql_command_result`, then, for a command only, one +//! `metrics` response (empty until command metrics exist; Spark sends it when `isCommand` and +//! skips it for a query), then `result_complete` for a reattachable operation: no schema or +//! batch response (Spark's `handleSqlCommand`). +//! +//! The statement is `SqlCommand.input` when set (what PySpark 4.x sends for every `spark.sql()`), +//! else the deprecated `sql` field. A command's result is one IPC stream; it is capped at +//! `0.7 x maxBatchSize` as Spark's single batch is (Spark asserts; Krishiv refuses with a typed +//! error). Unlike Spark's single batch it may hold several record batches. + +use std::sync::Arc; + +use krishiv_common::error::{KrishivError, KrishivUnsupportedFeature, OperationCanceled}; +use krishiv_connect::proto::v4_3 as pb; +use krishiv_connect::proto::v4_3::execute_plan_response::{ + Metrics, ResponseType, ResultComplete, SqlCommandResult, +}; +use krishiv_connect::proto::v4_3::plan::OpType; +use krishiv_connect::proto::v4_3::plan::compressed_operation::OpType as CompressedOpType; +use krishiv_connect::proto::v4_3::relation::RelType; +use krishiv_connect::translate::{sql_command_relation, sql_has_parameters}; +use krishiv_session::{QueryOptions, Session, SqlOutcome}; + +use super::{Sender, Stop, arrow, decode_command, decompress, settle}; +use crate::operation::ExecuteLimits; + +/// The `SqlCommand` of `plan`, when its command is one (a compressed command is decompressed +/// first, off the async workers, bounded by `maxPlanSize`). +pub(super) async fn of( + plan: &pb::Plan, + limits: &ExecuteLimits, +) -> Result, KrishivError> { + let command = match &plan.op_type { + Some(OpType::Command(command)) => command.command_type.clone(), + Some(OpType::CompressedOperation(c)) + if CompressedOpType::try_from(c.op_type) == Ok(CompressedOpType::Command) => + { + let c = (**c).clone(); + let max = limits.max_plan_size; + let bytes = tokio::task::spawn_blocking(move || decompress(&c, max)) + .await + .map_err(|e| -> KrishivError { + krishiv_common::error::InternalError { + message: format!("decompressing a plan: {e}"), + } + .into() + })??; + match &decode_command(&bytes)?.op_type { + Some(OpType::Command(command)) => command.command_type.clone(), + _ => None, + } + } + _ => None, + }; + Ok(match command { + Some(pb::command::CommandType::SqlCommand(sql)) => Some(*sql), + _ => None, + }) +} + +/// What a `sql_command` answers. +pub(super) struct Answer { + /// The relation of the `sql_command_result`. + pub(super) relation: pb::Relation, + /// Whether the statement was a command that ran (a `metrics` response follows). + pub(super) is_command: bool, +} + +/// Runs `command` on `session` under operation `id`: the relation to answer. The session time +/// zone of a command's result is read after the command ran (it may have been `SET`). +/// +/// # Errors +/// `KRISHIV_UNSUPPORTED_FEATURE` for an `input` that is not a plain `SQL` relation, for parameters +/// (the analyzer refuses SQL parameters on this base) and for a result over the batch cap; +/// `CONNECT_INVALID_PLAN.SQL_COMMAND_EXPECTS_SQL_OR_WITH_RELATIONS` for another `input`; the +/// statement's parse, analysis and run errors. +pub(super) async fn run( + session: &Session, + command: &pb::SqlCommand, + id: &str, + max_batch_size: usize, +) -> Result { + let relation = sql_command_relation(command)?; + let Some(RelType::Sql(sql)) = &relation.rel_type else { + return Err(internal( + "a SQL command did not translate to a SQL relation", + )); + }; + if sql_has_parameters(sql) { + return Err(KrishivUnsupportedFeature { + feature: "parameterised SQL (spark.connect.SqlCommand.args)".to_owned(), + } + .into()); + } + match session + .sql_command_with(&sql.query, QueryOptions::new().operation_id(id)) + .await? + { + SqlOutcome::Query(_) => Ok(Answer { + relation, + is_command: false, + }), + SqlOutcome::Command(df) => { + let tz: Option> = session + .conf() + .get("spark.sql.session.timeZone") + .map(Arc::from); + let tz = tz.as_ref(); + let schema = df.schema().await?; + let mut batches = Vec::new(); + for batch in df.collect().await? { + batches.push(arrow::spark_batch(&batch, tz)?); + } + if batches.is_empty() { + batches.push(arrow::empty_batch(&schema, tz)?); + } + let data = arrow::ipc_stream_of(&batches, None)?; + drop(batches); + #[expect( + clippy::cast_precision_loss, + clippy::cast_possible_truncation, + clippy::cast_sign_loss, + reason = "Spark's (maxBatchSize * 0.7).toLong on a positive size" + )] + let cap = (max_batch_size as f64 * 0.7) as usize; + if data.len() > cap { + return Err(KrishivUnsupportedFeature { + feature: format!( + "a SQL command result of {} bytes, over the {cap}-byte batch cap \ + (0.7 x spark.connect.grpc.arrow.maxBatchSize)", + data.len() + ), + } + .into()); + } + Ok(Answer { + relation: pb::Relation { + common: None, + rel_type: Some(RelType::LocalRelation(pb::LocalRelation { + data: Some(data), + schema: None, + })), + }, + is_command: true, + }) + } + _ => Err(internal("an unknown outcome of a SQL command")), + } +} + +fn internal(message: &str) -> KrishivError { + krishiv_common::error::InternalError { + message: message.to_owned(), + } + .into() +} + +/// Streams the answer of a `sql_command` and ends the operation. +pub(super) async fn encode(sender: Sender, answer: Answer) { + let op = Arc::clone(&sender.op); + let outcome = tokio::select! { + biased; + () = op.closed() => Err(Stop::Gone), + () = op.cancelled() => Err(Stop::Failed(OperationCanceled {}.into())), + r = produce(&sender, answer) => r, + }; + settle(&op, outcome); +} + +async fn produce(sender: &Sender, answer: Answer) -> Result<(), Stop> { + sender.admit().await?; + sender + .send( + sender.response(Some(ResponseType::SqlCommandResult(SqlCommandResult { + relation: Some(answer.relation), + }))), + ) + .await?; + if answer.is_command { + // The operator tree needs a session hook for a command's metrics: empty until then. + let mut metrics = sender.response(None); + metrics.metrics = Some(Metrics::default()); + sender.send(metrics).await?; + } + // As `produce` of the batches ends: atomic with interrupts. + if !sender.op.finish() { + return Err(Stop::Failed(OperationCanceled {}.into())); + } + if sender.op.reattachable() { + sender + .send(sender.response(Some(ResponseType::ResultComplete(ResultComplete {})))) + .await?; + } + sender.op.replay().finish(None); + Ok(()) +} diff --git a/crates/krishiv-connect-server/tests/execute.rs b/crates/krishiv-connect-server/tests/execute.rs index 0183d7ae6..928e1b837 100644 --- a/crates/krishiv-connect-server/tests/execute.rs +++ b/crates/krishiv-connect-server/tests/execute.rs @@ -529,7 +529,9 @@ async fn compressed_plans_decode_and_a_bomb_fails_before_it_expands() { e.message() ); let command = pb::Command { - command_type: Some(pb::command::CommandType::SqlCommand(Box::default())), + command_type: Some(pb::command::CommandType::GetResourcesCommand(Box::new( + pb::GetResourcesCommand {}, + ))), }; let data = zstd::bulk::compress(&command.encode_to_vec(), 3).expect("zstd"); let cmd = pb::plan::compressed_operation::OpType::Command as i32; @@ -571,7 +573,9 @@ async fn invalid_requests_are_refused_with_sparks_classes() { ); let command = request(pb::Plan { op_type: Some(OpType::Command(Box::new(pb::Command { - command_type: Some(pb::command::CommandType::SqlCommand(Box::default())), + command_type: Some(pb::command::CommandType::GetResourcesCommand(Box::new( + pb::GetResourcesCommand {}, + ))), }))), }); let e = collect(&run.channel, command).await.unwrap_err(); @@ -1032,3 +1036,521 @@ async fn large_row_result() { .expect("utf8"); assert_eq!(s.value(0).len(), 65_536); } + +// ---- local relations, schemas, qualifiers and expression strings (P3-T012a part 2b) -------- + +/// An Arrow IPC stream of `(id: Int64, name: Utf8)` rows, as PySpark's `createDataFrame` sends. +fn people_ipc() -> Vec { + use arrow_array::{Int64Array, StringArray}; + use arrow_schema::{DataType, Field, Schema}; + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, true), + Field::new("name", DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int64Array::from(vec![1, 2, 3])), + Arc::new(StringArray::from(vec![Some("a"), None, Some("c")])), + ], + ) + .expect("batch"); + let mut out = Vec::new(); + { + let mut w = arrow_ipc::writer::StreamWriter::try_new(&mut out, &schema).expect("writer"); + w.write(&batch).expect("write"); + w.finish().expect("finish"); + } + out +} + +fn local(data: Option>, schema: Option<&str>) -> pb::Relation { + relation(RelType::LocalRelation(pb::LocalRelation { + data, + schema: schema.map(str::to_owned), + })) +} + +/// The one Arrow batch of a small result. +fn rows(responses: &[pb::ExecutePlanResponse]) -> RecordBatch { + let mut all: Vec = batches(responses) + .iter() + .flat_map(|b| read_ipc(&b.data)) + .collect(); + assert_eq!(all.len(), 1, "a small result is one batch"); + all.remove(0) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn a_local_relation_is_collected() { + let run = start().await; + let r = collect(&run.channel, request(root(local(Some(people_ipc()), None)))) + .await + .expect("runs"); + let got = rows(&r); + assert_eq!(got.num_rows(), 3); + assert_eq!(got.schema().field(0).name(), "id"); + assert_eq!(got.schema().field(1).name(), "name"); + // The batches live once, in the session's store, held by the plan's lease. + let again = collect(&run.channel, request(root(local(Some(people_ipc()), None)))) + .await + .expect("runs again"); + assert_eq!(rows(&again).num_rows(), 3); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn a_local_relation_applies_its_schema_eagerly() { + let run = start().await; + // The user schema renames by position and casts `id` to int. + let r = collect( + &run.channel, + request(root(local(Some(people_ipc()), Some("a INT, b STRING")))), + ) + .await + .expect("runs"); + let got = rows(&r); + assert_eq!(got.schema().field(0).name(), "a"); + assert_eq!( + got.schema().field(0).data_type(), + &arrow_schema::DataType::Int32 + ); + assert_eq!(got.num_rows(), 3); + // A schema of a different column count fails while the plan is built. + let e = collect( + &run.channel, + request(root(local(Some(people_ipc()), Some("a INT")))), + ) + .await + .unwrap_err(); + assert_eq!(error_class(&e), "NUM_COLUMNS_MISMATCH"); + // A schema alone is an empty relation; neither is invalid. + let empty = collect( + &run.channel, + request(root(local(None, Some("a INT, b STRING")))), + ) + .await + .expect("runs"); + assert_eq!(rows(&empty).num_rows(), 0); + let e = collect(&run.channel, request(root(local(None, None)))) + .await + .unwrap_err(); + assert_eq!( + error_class(&e), + "CONNECT_INVALID_PLAN.SCHEMA_REQUIRED_FOR_LOCAL_RELATION" + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn a_cached_local_relation_without_its_artifact_is_not_found() { + let run = start().await; + let cached = relation(RelType::CachedLocalRelation(pb::CachedLocalRelation { + hash: "0123abcd".to_owned(), + })); + let e = collect(&run.channel, request(root(cached))) + .await + .unwrap_err(); + assert_eq!( + error_class(&e), + "CONNECT_INVALID_PLAN.NOT_FOUND_CACHED_LOCAL_RELATION" + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn two_local_relations_join_in_one_plan() { + let run = start().await; + let join = relation(RelType::Join(Box::new(pb::Join { + left: Some(Box::new(local(Some(people_ipc()), None))), + right: Some(Box::new(local( + Some(people_ipc()), + Some("id2 BIGINT, n2 STRING"), + ))), + join_type: pb::join::JoinType::Cross as i32, + ..Default::default() + }))); + let r = collect(&run.channel, request(root(join))) + .await + .expect("runs"); + assert_eq!(rows(&r).num_rows(), 9); +} + +fn column(name: &str) -> pb::Expression { + pb::Expression { + common: None, + expr_type: Some(pb::expression::ExprType::UnresolvedAttribute( + pb::expression::UnresolvedAttribute { + unparsed_identifier: name.to_owned(), + plan_id: None, + is_metadata_column: None, + }, + )), + } +} + +fn expression_string(text: &str) -> pb::Expression { + pb::Expression { + common: None, + expr_type: Some(pb::expression::ExprType::ExpressionString( + pb::expression::ExpressionString { + expression: text.to_owned(), + }, + )), + } +} + +fn project(input: pb::Relation, exprs: Vec) -> pb::Relation { + relation(RelType::Project(Box::new(pb::Project { + input: Some(Box::new(input)), + expressions: exprs, + }))) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn an_expression_string_runs_as_select_expr() { + let run = start().await; + let plan = project( + local(Some(people_ipc()), None), + vec![expression_string("id * 10 AS tens")], + ); + let r = collect(&run.channel, request(root(plan))) + .await + .expect("runs"); + let got = rows(&r); + assert_eq!(got.schema().field(0).name(), "tens"); + let values = got + .column(0) + .as_any() + .downcast_ref::() + .expect("longs"); + assert_eq!(values.values().to_vec(), vec![10, 20, 30]); + let e = collect( + &run.channel, + request(root(project( + local(Some(people_ipc()), None), + vec![expression_string("id +")], + ))), + ) + .await + .unwrap_err(); + assert_eq!(error_class(&e), "PARSE_SYNTAX_ERROR"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn a_subquery_alias_qualifier_qualifies_the_columns() { + let run = start().await; + let aliased = |qualifier: &[&str]| { + relation(RelType::SubqueryAlias(Box::new(pb::SubqueryAlias { + input: Some(Box::new(local(Some(people_ipc()), None))), + alias: "t".to_owned(), + qualifier: qualifier.iter().map(|q| (*q).to_owned()).collect(), + }))) + }; + // `db.t.id` and `t.id` both name the column of `db.t`. + for column_name in ["db.t.id", "t.id", "id"] { + let r = collect( + &run.channel, + request(root(project(aliased(&["db"]), vec![column(column_name)]))), + ) + .await + .unwrap_or_else(|e| panic!("{column_name}: {e}")); + assert_eq!(rows(&r).num_rows(), 3, "{column_name}"); + } + // The qualifier is not a prefix of another relation's name. + let e = collect( + &run.channel, + request(root(project(aliased(&["db"]), vec![column("other.t.id")]))), + ) + .await + .unwrap_err(); + assert_eq!(error_class(&e), "UNRESOLVED_COLUMN.WITH_SUGGESTION"); + // Without a qualifier `db.t.id` does not resolve. + let e = collect( + &run.channel, + request(root(project(aliased(&[]), vec![column("db.t.id")]))), + ) + .await + .unwrap_err(); + assert!( + error_class(&e).starts_with("UNRESOLVED_COLUMN"), + "{}", + error_class(&e) + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn a_data_source_schema_declares_the_columns() { + let dir = tempfile::tempdir().expect("dir"); + let path = dir.path().join("t.csv"); + std::fs::write(&path, "1,x\n2,y\n").expect("write"); + let read = relation(RelType::Read(pb::Read { + is_streaming: false, + read_type: Some(pb::read::ReadType::DataSource(pb::read::DataSource { + format: Some("csv".to_owned()), + schema: Some("n INT, s STRING".to_owned()), + paths: vec![path.to_string_lossy().into_owned()], + ..Default::default() + })), + })); + let run = start().await; + let r = collect(&run.channel, request(root(read))) + .await + .expect("runs"); + let got = rows(&r); + assert_eq!(got.schema().field(0).name(), "n"); + assert_eq!( + got.schema().field(0).data_type(), + &arrow_schema::DataType::Int32 + ); + assert_eq!(got.num_rows(), 2); +} + +fn sql_command(query: &str) -> pb::Plan { + #[expect(deprecated, reason = "the field every client version sends")] + let command = pb::SqlCommand { + sql: query.to_owned(), + ..Default::default() + }; + pb::Plan { + op_type: Some(OpType::Command(Box::new(pb::Command { + command_type: Some(pb::command::CommandType::SqlCommand(Box::new(command))), + }))), + } +} + +/// The command PySpark 4.x sends for `spark.sql(query)`: the statement as `input`, a `SQL` +/// relation (the deprecated fields unset). +fn sql_command_input(query: &str, plan_id: i64) -> pb::Plan { + let command = pb::SqlCommand { + input: Some(pb::Relation { + common: Some(pb::RelationCommon { + plan_id: Some(plan_id), + ..Default::default() + }), + rel_type: Some(RelType::Sql(pb::Sql { + query: query.to_owned(), + ..Default::default() + })), + }), + ..Default::default() + }; + pb::Plan { + op_type: Some(OpType::Command(Box::new(pb::Command { + command_type: Some(pb::command::CommandType::SqlCommand(Box::new(command))), + }))), + } +} + +/// The relation of the one `sql_command_result` response. +fn sql_result(responses: &[pb::ExecutePlanResponse]) -> pb::Relation { + let results: Vec<_> = responses + .iter() + .filter_map(|r| match &r.response_type { + Some(ResponseType::SqlCommandResult(c)) => c.relation.clone(), + _ => None, + }) + .collect(); + assert_eq!(results.len(), 1, "one sql_command_result"); + assert!( + batches(responses).is_empty(), + "a SQL command streams no batch" + ); + results.into_iter().next().expect("one") +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn a_sql_query_answers_the_sql_relation_the_client_sends_again() { + let run = start().await; + let r = collect( + &run.channel, + reattachable(request(sql_command("SELECT 1 AS a, 'x' AS b"))), + ) + .await + .expect("runs"); + // Only the result and `result_complete`: no schema, batch or metrics. + assert_eq!(r.len(), 2, "{r:?}"); + assert!(matches!( + r[1].response_type, + Some(ResponseType::ResultComplete(_)) + )); + let relation = sql_result(&r); + let Some(RelType::Sql(sql)) = &relation.rel_type else { + panic!("a query answers a SQL relation: {relation:?}"); + }; + assert_eq!(sql.query, "SELECT 1 AS a, 'x' AS b"); + // The client sends it as a root relation. + let again = collect(&run.channel, request(root(relation))) + .await + .expect("runs"); + let got = rows(&again); + assert_eq!(got.num_rows(), 1); + assert_eq!(got.schema().field(0).name(), "a"); + // An analysis error is raised by `spark.sql` itself, as Spark does. + let e = collect( + &run.channel, + request(sql_command("SELECT * FROM no_such_table")), + ) + .await + .unwrap_err(); + assert_eq!(error_class(&e), "TABLE_OR_VIEW_NOT_FOUND"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn a_sql_command_runs_at_once_and_answers_its_rows_as_a_local_relation() { + let run = start().await; + let r = collect( + &run.channel, + request(sql_command("CREATE TEMPORARY VIEW v AS SELECT 7 AS n")), + ) + .await + .expect("runs"); + // Spark's order for a command: the result, then metrics (then `result_complete`). + assert!( + matches!(r[0].response_type, Some(ResponseType::SqlCommandResult(_))), + "{r:?}" + ); + assert!(r[1].metrics.is_some(), "{r:?}"); + let relation = sql_result(&r); + assert!( + matches!(relation.rel_type, Some(RelType::LocalRelation(_))), + "{relation:?}" + ); + // It ran: the view is there for the next request, and no later request creates it again. + let entry = run + .manager + .get(&Principal::local_user(), SESSION, None) + .expect("session"); + let df = entry + .session() + .sql("SELECT n FROM v") + .await + .expect("the view exists"); + assert_eq!(df.collect().await.expect("rows").len(), 1); + // `SET key=value` returns its rows (key, value) as a local relation the client re-sends. + let r = collect( + &run.channel, + request(sql_command("SET spark.sql.shuffle.partitions=3")), + ) + .await + .expect("runs"); + let relation = sql_result(&r); + let again = collect(&run.channel, request(root(relation))) + .await + .expect("a command's relation runs"); + let got = rows(&again); + assert_eq!(got.num_rows(), 1); + assert_eq!(got.num_columns(), 2); + assert_eq!( + entry.session().conf().get("spark.sql.shuffle.partitions"), + Some("3".to_owned()) + ); + // A failing command is a failed start. + let e = collect(&run.channel, request(sql_command("DROP VIEW nope_view"))) + .await + .unwrap_err(); + assert_eq!(error_class(&e), "TABLE_OR_VIEW_NOT_FOUND"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn a_compressed_sql_command_runs_and_unsupported_forms_are_refused() { + let run = start().await; + #[expect(deprecated, reason = "the field every client version sends")] + let command = pb::Command { + command_type: Some(pb::command::CommandType::SqlCommand(Box::new( + pb::SqlCommand { + sql: "SELECT 2".to_owned(), + ..Default::default() + }, + ))), + }; + let data = zstd::bulk::compress(&command.encode_to_vec(), 3).expect("zstd"); + let plan = pb::Plan { + op_type: Some(OpType::CompressedOperation(Box::new( + pb::plan::CompressedOperation { + data, + op_type: pb::plan::compressed_operation::OpType::Command as i32, + compression_codec: pb::CompressionCodec::Zstd as i32, + }, + ))), + }; + let r = collect(&run.channel, request(plan)).await.expect("runs"); + assert!(matches!(sql_result(&r).rel_type, Some(RelType::Sql(_)))); + // Parameters are refused up front (golden-pending: group 4). + let mut with_args = sql_command("SELECT ?"); + if let Some(OpType::Command(c)) = &mut with_args.op_type + && let Some(pb::command::CommandType::SqlCommand(s)) = &mut c.command_type + { + #[expect(deprecated, reason = "the field every client version sends")] + { + s.pos_arguments = vec![column("x")]; + } + } + let e = collect(&run.channel, request(with_args)).await.unwrap_err(); + assert_eq!(error_class(&e), "KRISHIV_UNSUPPORTED_FEATURE"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn a_pyspark_4_sql_command_sends_its_statement_as_input() { + let run = start().await; + // A query: the `input` relation comes back as it was, plan id included. + let r = collect( + &run.channel, + reattachable(request(sql_command_input("SELECT 1 AS a", 41))), + ) + .await + .expect("runs"); + assert_eq!(r.len(), 2, "{r:?}"); + let relation = sql_result(&r); + assert_eq!(relation.common.as_ref().and_then(|c| c.plan_id), Some(41)); + assert!( + matches!(&relation.rel_type, Some(RelType::Sql(s)) if s.query == "SELECT 1 AS a"), + "{relation:?}" + ); + // A command through `input` runs at once. + let r = collect( + &run.channel, + request(sql_command_input( + "CREATE TEMPORARY VIEW iv AS SELECT 5 AS n", + 42, + )), + ) + .await + .expect("runs"); + assert!(matches!( + sql_result(&r).rel_type, + Some(RelType::LocalRelation(_)) + )); + let entry = run + .manager + .get(&Principal::local_user(), SESSION, None) + .expect("session"); + let df = entry + .session() + .sql("SELECT n FROM iv") + .await + .expect("the view exists"); + assert_eq!(df.collect().await.expect("rows").len(), 1); + // Another input relation is refused with Spark's class; WithRelations is group 4. + let mut other = sql_command_input("x", 1); + if let Some(OpType::Command(c)) = &mut other.op_type + && let Some(pb::command::CommandType::SqlCommand(s)) = &mut c.command_type + { + s.input = Some(range(1)); + } + let e = collect(&run.channel, request(other)).await.unwrap_err(); + assert_eq!( + error_class(&e), + "CONNECT_INVALID_PLAN.SQL_COMMAND_EXPECTS_SQL_OR_WITH_RELATIONS" + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn a_command_result_over_the_batch_cap_is_refused() { + let run = start_with(ExecuteLimits { + max_batch_size: 64, + ..ExecuteLimits::default() + }) + .await; + let e = collect(&run.channel, request(sql_command("SHOW TABLES"))) + .await + .unwrap_err(); + assert_eq!(error_class(&e), "KRISHIV_UNSUPPORTED_FEATURE"); +} diff --git a/crates/krishiv-connect-server/tests/rpcs.rs b/crates/krishiv-connect-server/tests/rpcs.rs index 900ea20bb..8752c73c0 100644 --- a/crates/krishiv-connect-server/tests/rpcs.rs +++ b/crates/krishiv-connect-server/tests/rpcs.rs @@ -311,6 +311,25 @@ async fn analyze_explain_tree_string_and_flags() { panic!("is local"); }; assert!(!l.is_local); + // A `LocalRelation` (`createDataFrame`) is local; a projection over one is not. + let local = pb::Plan { + op_type: Some(PlanOp::Root(Box::new(pb::Relation { + common: None, + rel_type: Some(RelType::LocalRelation(pb::LocalRelation { + data: None, + schema: Some("a INT".to_owned()), + })), + }))), + }; + let AResult::IsLocal(l) = analyzed( + &run, + Analyze::IsLocal(Box::new(areq::IsLocal { plan: Some(local) })), + ) + .await + else { + panic!("is local"); + }; + assert!(l.is_local); // Codex review (#738): SQL text is parsed during analysis, so a SQL // `VALUES` is not an inline table when `is_local` asks (Spark's answer // for it is to be checked against 4.2/4.3 before this changes). diff --git a/crates/krishiv-connect/CLAUDE.md b/crates/krishiv-connect/CLAUDE.md index 135074a28..fc80931d0 100644 --- a/crates/krishiv-connect/CLAUDE.md +++ b/crates/krishiv-connect/CLAUDE.md @@ -11,8 +11,8 @@ the logical-plan IR of `krishiv-plan`. The server (P3-T011, P3-T012) extends thi - `src/wire.rs` the only decode path (`wire::decode_*`, `decode_*_with(bytes, &DecodeOptions)`; `decode_request` for any RPC request, the `Request` trait the generated layout implements): size, count, depth (`krishiv.connect.maxPlanDepth`) and field-version checks - `src/proto/mod.rs` versioned generated protos; `src/proto/v4_3/spark.connect.pb.rs` generated messages (4.2 is kept only for the golden descriptor) (do not edit) -- `src/translate/mod.rs` `translate_*` entry points, `command_not_supported` (commands are 56K00 until T012); `src/translate/relation.rs` and `expression.rs` (IR building) -- `src/translate/literal.rs` literals; `src/translate/names.rs` names and plan-id origins +- `src/translate/mod.rs` `translate_*` entry points, `command_not_supported` (the server runs `sql_command`; every other command is 56K00 until T012); `src/translate/relation.rs` and `expression.rs` (IR building) +- `src/translate/literal.rs` literals (scalars, intervals, nested array/map/struct, specialized arrays); `src/translate/local.rs` the host protocol for local relations; `src/translate/schema.rs` schema texts; `src/translate/names.rs` names and plan-id origins - `src/translate/errors.rs` `KRISHIV_UNSUPPORTED_FEATURE` and `CONNECT_INVALID_PLAN.*` errors - `src/arrow.rs` result columns Spark sends in another Arrow shape than the engine holds (geometry/geography) - `tests/fuzz_regressions.rs` fuzz findings (inputs that once ran out of memory) must end in a typed error @@ -47,9 +47,19 @@ the logical-plan IR of `krishiv-plan`. The server (P3-T011, P3-T012) extends thi text, literals, names, options, paths), only their sizes. `Decoded` wipes option values, paths and `Any` payloads with zeros before it drops the message (`Scrub`, `scrub.pb.rs`); a proto built in Rust and passed to `translate_*` stays the caller's to wipe. -- SQL expression strings (`selectExpr`, `where("...")`, `expr()`) are unsupported until the analyzer lowers the SQL - parser's syntax tree into the IR (P2-T071a); `ToDF`, pivots, user schemas and nested-type literals wait for - IR nodes or later groups (`src/coverage.rs` lists every arm). +- SQL expression strings (`selectExpr`, `where("...")`, `expr()`) are parsed eagerly while translating, by + `krishiv_sql::parse_and_lower_named_expression` with the session's `TranslateConf::parser` (filled by + `krishiv-connect-server` from `Session::parser_options`); `ToDF`, pivots and `TIME` literals wait for IR nodes + or later groups (`src/coverage.rs` lists every arm). +- **Local relations** (`createDataFrame`): the translator is synchronous, so the host builds them. It translates the + plan once with `LocalRelations::collecting()` (records a `LocalRequest` per `local_relation` / + `cached_local_relation`, answers with placeholders), builds the relations (`Session::local_relation`), then + translates again with `LocalRelations::resolved(..)`: the n-th request gets the n-th relation. A plan without + local relations is translated once. `TranslateConf::default()` has no host: a local relation is + `KRISHIV_UNSUPPORTED_FEATURE`. `chunked_cached_local_relation` stays 56K00. +- Schema texts (`translate/schema.rs`): DDL, else JSON (`parseTypeWithFallback`); a local relation's non-struct type is + wrapped as `struct`, a data source's must be a struct (`CONNECT_INVALID_PLAN.INVALID_SCHEMA_NON_STRUCT_TYPE`). + `SubqueryAlias.qualifier ++ [alias]` is one `ALIAS_PART_SEP`-joined alias. - A `oneof` arm that is not translated (`Relation`, `Expression`, `Literal`, `DataType`, later `Command` and `Catalog`) fails with Spark's `CONNECT_INVALID_PLAN.INVALID_ONE_OF_FIELD_NOT_SUPPORTED` (SQLSTATE 56K00; parameters fullName, the upper-cased field name and the field number, as `InvalidInputErrors.invalidOneOfField`; owner decision 6, diff --git a/crates/krishiv-connect/Cargo.toml b/crates/krishiv-connect/Cargo.toml index 003fa86ed..80ff3c8e9 100644 --- a/crates/krishiv-connect/Cargo.toml +++ b/crates/krishiv-connect/Cargo.toml @@ -16,6 +16,7 @@ arrow-array.workspace = true arrow-schema.workspace = true krishiv-common = { path = "../krishiv-common", version = "=0.2.0" } krishiv-plan = { path = "../krishiv-plan", version = "=0.2.0" } +krishiv-sql = { path = "../krishiv-sql", version = "=0.2.0" } krishiv-types = { path = "../krishiv-types", version = "=0.2.0" } # `no-recursion-limit`: prost's fixed limit of 100 nested messages rejects # ordinary plans (a loop of `withColumn` calls); `wire::decode_plan` bounds the diff --git a/crates/krishiv-connect/src/coverage.rs b/crates/krishiv-connect/src/coverage.rs index 3a2168f50..db1b12b26 100644 --- a/crates/krishiv-connect/src/coverage.rs +++ b/crates/krishiv-connect/src/coverage.rs @@ -96,7 +96,7 @@ pub fn all() -> Vec { } oneof! { relation, RELATION_ROWS, relation::RelType, "Relation.rel_type"; - Read => ("read", 2, Partial, "DataSource.schema, predicates (JDBC) and is_streaming", "Read(ReadSource::{Table,Files}); user schema on ReadSource::Files; predicates → JDBC partitions (refuse as Spark does for non-JDBC formats: PREDICATES_NOT_SUPPORTED_FOR_DATA_SOURCE)", "P3-T012", 1), + Read => ("read", 2, Partial, "DataSource.predicates (JDBC) and is_streaming", "Read(ReadSource::{Table,Files}); DataSource.schema → ReadSource::Files schema (DDL or JSON, a non-struct is INVALID_SCHEMA_NON_STRUCT_TYPE); predicates → JDBC partitions (refuse as Spark does for non-JDBC formats: PREDICATES_NOT_SUPPORTED_FOR_DATA_SOURCE)", "P3-T012", 1), Project => ("project", 3, Done, "", "Project; missing input → OneRowRelation", "P3-T012", 1), Filter => ("filter", 4, Done, "", "Filter", "P2-T091a", 0), Join => ("join", 5, Partial, "join_data_type (joinWith, typed)", "Join; join_data_type is Scala joinWith only → unsupported row", "P3-T012", 1), @@ -105,12 +105,12 @@ oneof! { relation, RELATION_ROWS, relation::RelType, "Relation.rel_type"; Limit => ("limit", 8, Done, "", "Limit", "P2-T091a", 0), Aggregate => ("aggregate", 9, Partial, "GROUP_TYPE_PIVOT", "Aggregate; pivot → Pivot (exists)", "P3-T012", 2), Sql => ("sql", 10, Partial, "DataFrame arguments (with_relations)", "Sql; the with_relations form (#41) for DataFrame arguments", "P3-T012", 4), - LocalRelation => ("local_relation", 11, Missing, "", "InlineTable from the Arrow IPC data plus the DDL or JSON schema (createDataFrame); a size limit on the data", "P3-T012", 1), + LocalRelation => ("local_relation", 11, Done, "", "Read(ReadSource::Local) from the Arrow IPC data plus the DDL or JSON schema (createDataFrame), built by the host before translation (TranslateConf::local_relations); no data and no schema is SCHEMA_REQUIRED_FOR_LOCAL_RELATION; the inline size is bounded by the gRPC message limit", "P3-T012a", 1), Sample => ("sample", 12, Missing, "", "new Sample (Bernoulli/Poisson with seed, Spark's XORShiftRandom order)", "P3-T012", 3), Offset => ("offset", 13, Done, "", "Offset", "P2-T091a", 0), Deduplicate => ("deduplicate", 14, Missing, "", "lowering: Distinct (all columns) or Aggregate with first per non-key column; within_watermark → P4", "P3-T012", 3), Range => ("range", 15, Done, "", "Range", "P2-T091a", 0), - SubqueryAlias => ("subquery_alias", 16, Partial, "qualifier", "SubqueryAlias with the qualifier", "P3-T012", 1), + SubqueryAlias => ("subquery_alias", 16, Done, "", "SubqueryAlias; the qualifier and the alias are the alias parts joined by ALIAS_PART_SEP", "P3-T012", 1), Repartition => ("repartition", 17, Missing, "", "Repartition (round robin / coalesce)", "P3-T012", 3), ToDf => ("to_df", 18, Missing, "", "lowering: Project of Alias (positional)", "P3-T012", 4), WithColumnsRenamed => ("with_columns_renamed", 19, Missing, "", "lowering: Project", "P3-T012", 4), @@ -130,7 +130,7 @@ oneof! { relation, RELATION_ROWS, relation::RelType, "Relation.rel_type"; WithWatermark => ("with_watermark", 33, OutOfScope, "", "P4", "", 0), ApplyInPandasWithState => ("apply_in_pandas_with_state", 34, OutOfScope, "", "P4", "", 0), HtmlString => ("html_string", 35, Missing, "", "as show_string, HTML form", "P3-T012", 4), - CachedLocalRelation => ("cached_local_relation", 36, Missing, "", "InlineTable from the artifact cache (AddArtifacts cache/, T011); errors NOT_FOUND_CACHED_LOCAL_RELATION", "P3-T083", 7), + CachedLocalRelation => ("cached_local_relation", 36, Partial, "hashes in the artifact store (AddArtifacts cache/ is not stored yet: every hash is NOT_FOUND_CACHED_LOCAL_RELATION)", "Read(ReadSource::Local) through the same host request as local_relation; the artifact lookup is the host's", "P3-T083", 7), CachedRemoteRelation => ("cached_remote_relation", 37, Missing, "", "the session's cached-DataFrame registry (T083)", "P3-T083", 7), CommonInlineUserDefinedTableFunction => ("common_inline_user_defined_table_function", 38, Missing, "", "TableFunction with a Python UDTF (T014)", "P3-T014", 6), AsOfJoin => ("as_of_join", 39, Missing, "", "new AsOfJoin, shared with the ASOF JOIN SQL syntax", "P3-T012", 8), @@ -161,10 +161,10 @@ oneof! { relation, RELATION_ROWS, relation::RelType, "Relation.rel_type"; } oneof! { expression, EXPRESSION_ROWS, expression::ExprType, "Expression.expr_type"; - Literal => ("literal", 1, Partial, "Array, Map, Struct, SpecializedArray, Time", "Literal/TypedLiteral; nested literals after the IR literal supports them; nanosecond timestamps are refused until the types exist", "P3-T012", 1), + Literal => ("literal", 1, Partial, "Time", "Literal/TypedLiteral; nested literals are Literal::{Array,KeyValueMap,Struct}; nanosecond timestamps are refused until the types exist", "P3-T012", 1), UnresolvedAttribute => ("unresolved_attribute", 2, Partial, "is_metadata_column", "Column; metadata columns (_metadata) → the scan's metadata columns", "P3-T012", 1), UnresolvedFunction => ("unresolved_function", 3, Partial, "is_user_defined_function (qualified names) and is_internal", "Function with qualifier; internal functions map to Spark's internal names where Krishiv has them, otherwise 56K00-style refusal naming the function", "P3-T012", 1), - ExpressionString => ("expression_string", 4, Missing, "", "the SQL expression parser + analyzer lowering that select_expr/expr() use (krishiv-session)", "P3-T012", 1), + ExpressionString => ("expression_string", 4, Done, "", "krishiv_sql::parse_and_lower_named_expression, as Session::select_expr and Session::expr (parser options from TranslateConf)", "P3-T012a", 1), UnresolvedStar => ("unresolved_star", 5, Done, "", "Star", "P2-T091a", 0), Alias => ("alias", 6, Partial, "metadata and several names", "Alias; several names → MultiAlias (exists)", "P3-T012", 1), Cast => ("cast", 7, Done, "", "Cast", "P2-T091a", 0), @@ -190,7 +190,7 @@ oneof! { command, COMMAND_ROWS, command::CommandType, "Command.command_type"; WriteOperation => ("write_operation", 2, Missing, "", "DataFrameWriter V1: krishiv-session save/save_as_table/insert_into (dataframe.rs), semantics T023", "P3-T023", 11), CreateDataframeView => ("create_dataframe_view", 3, Missing, "", "create_or_replace_temp_view (exists), global temp views", "P3-T012", 4), WriteOperationV2 => ("write_operation_v2", 4, Missing, "", "DataFrameWriterV2 (create/replace/append/overwrite/overwrite_partitions exist in dataframe.rs), T023", "P3-T023", 11), - SqlCommand => ("sql_command", 5, Missing, "", "spark.sql() sends this. Run commands at once; for queries return sql_command_result.relation (Spark returns a SQL or LocalRelation relation the client re-sends). Without it no spark.sql test passes", "P3-T012", 1), + SqlCommand => ("sql_command", 5, Partial, "input of WithRelations form (group 4; the plain SQL input is supported); named and positional arguments (the analyzer refuses SQL parameters)", "spark.sql() sends this. Run commands at once; for queries return sql_command_result.relation (Spark returns a SQL or LocalRelation relation the client re-sends). Translated here as the SQL relation of a query; the server runs commands (execute.rs)", "P3-T012", 1), WriteStreamOperationStart => ("write_stream_operation_start", 6, OutOfScope, "", "P4", "", 0), StreamingQueryCommand => ("streaming_query_command", 7, OutOfScope, "", "P4", "", 0), GetResourcesCommand => ("get_resources_command", 8, Missing, "", "an empty resource map (single node, no executors)", "P3-T012", 4), @@ -265,11 +265,11 @@ oneof! { literal, LITERAL_ROWS, expression::literal::LiteralType, "Expression.Li CalendarInterval => ("calendar_interval", 19, Done, "", "Literal / TypedLiteral", "P2-T091a", 1), YearMonthInterval => ("year_month_interval", 20, Done, "", "Literal / TypedLiteral", "P2-T091a", 1), DayTimeInterval => ("day_time_interval", 21, Done, "", "Literal / TypedLiteral", "P2-T091a", 1), - Array => ("array", 22, Missing, "", "nested and TIME literals wait for IR values", "P3-T012", 1), - Map => ("map", 23, Missing, "", "nested and TIME literals wait for IR values", "P3-T012", 1), - Struct => ("struct", 24, Missing, "", "nested and TIME literals wait for IR values", "P3-T012", 1), - SpecializedArray => ("specialized_array", 25, Missing, "", "nested and TIME literals wait for IR values", "P3-T012", 1), - Time => ("time", 26, Missing, "", "nested and TIME literals wait for IR values", "P3-T012", 1), + Array => ("array", 22, Done, "", "Literal::{Array,KeyValueMap,Struct}; the type is the root data_type, else the message\'s own type fields", "P3-T012a", 1), + Map => ("map", 23, Done, "", "Literal::{Array,KeyValueMap,Struct}; the type is the root data_type, else the message\'s own type fields", "P3-T012a", 1), + Struct => ("struct", 24, Done, "", "Literal::{Array,KeyValueMap,Struct}; the type is the root data_type, else the message\'s own type fields", "P3-T012a", 1), + SpecializedArray => ("specialized_array", 25, Done, "", "Literal::{Array,KeyValueMap,Struct}; the type is the root data_type, else the message\'s own type fields", "P3-T012a", 1), + Time => ("time", 26, Missing, "", "TIME literals wait for an IR value", "P3-T012", 1), TimestampNtzNanos => ("timestamp_ntz_nanos", 29, Missing, "", "no IR type for nanosecond timestamps", "P3-T012", 1), TimestampLtzNanos => ("timestamp_ltz_nanos", 30, Missing, "", "no IR type for nanosecond timestamps", "P3-T012", 1), } diff --git a/crates/krishiv-connect/src/translate/expression.rs b/crates/krishiv-connect/src/translate/expression.rs index 3f4a1a674..303ce839d 100644 --- a/crates/krishiv-connect/src/translate/expression.rs +++ b/crates/krishiv-connect/src/translate/expression.rs @@ -89,6 +89,18 @@ impl Translator<'_> { )); } ExprType::UnresolvedStar(s) => ExprNode::Star(star(s)?), + ExprType::ExpressionString(x) => { + // `expr()` and `selectExpr`: the SQL expression parser and lowering the session's + // own `select_expr` uses, with the session's parser options, eagerly (Spark's + // `parser.parseExpression` at translation). Spans index no registered source, as + // for the DataFrame API. + return krishiv_sql::parse_and_lower_named_expression( + &x.expression, + &self.conf.parser, + krishiv_plan::SourceId(u32::MAX), + ) + .map_err(|e| e.error); + } ExprType::Alias(a) => { if a.metadata.as_deref().is_some_and(|m| !m.is_empty()) { return Err(unsupported_field( diff --git a/crates/krishiv-connect/src/translate/expression/tests.rs b/crates/krishiv-connect/src/translate/expression/tests.rs index 16a044dcd..8b477fa70 100644 --- a/crates/krishiv-connect/src/translate/expression/tests.rs +++ b/crates/krishiv-connect/src/translate/expression/tests.rs @@ -497,25 +497,6 @@ fn stars() { ); } -#[test] -fn sql_expression_strings_are_unsupported_for_now() { - let sql = expr(ExprType::ExpressionString(expression::ExpressionString { - expression: "a > 1".into(), - })); - let e = err(&sql); - assert_eq!( - e.error_class(), - ErrorClass::ConnectInvalidPlanInvalidOneOfFieldNotSupported - ); - assert!( - e.message().contains("EXPRESSION_STRING(4)"), - "{}", - e.message() - ); - // The expression text is not echoed into the error. - assert!(!e.message().contains("a > 1")); -} - #[test] fn unsupported_expression_kinds_are_named() { for (kind, name) in [ @@ -671,3 +652,43 @@ fn on_small_stack(stack: usize, f: impl FnOnce() + Send + 'static) { .join() .unwrap(); } + +#[test] +fn an_expression_string_is_parsed_like_select_expr() { + let e = expr(ExprType::ExpressionString(expression::ExpressionString { + expression: "a + 1 AS b".to_owned(), + })); + assert_eq!(text(&e), text_of_sql("a + 1 AS b")); + // The parser's own error, with its class, and no echo of the text. + let bad = expr(ExprType::ExpressionString(expression::ExpressionString { + expression: "CANARY +".to_owned(), + })); + let e = err(&bad); + assert_eq!( + e.error_class(), + ErrorClass::from_name("PARSE_SYNTAX_ERROR").unwrap() + ); +} + +#[test] +fn an_expression_string_is_parsed_with_the_conf_parser_options() { + // The options are the conf's: the same text under the default options translates. + let e = expr(ExprType::ExpressionString(expression::ExpressionString { + expression: "a > 1".to_owned(), + })); + let conf = TranslateConf { + parser: krishiv_sql::ParserOptions::default(), + ..TranslateConf::default() + }; + assert!(translate_expression(&e, &conf).is_ok()); +} + +fn text_of_sql(sql: &str) -> String { + krishiv_sql::parse_and_lower_named_expression( + sql, + &krishiv_sql::ParserOptions::default(), + krishiv_plan::SourceId(u32::MAX), + ) + .unwrap() + .to_string() +} diff --git a/crates/krishiv-connect/src/translate/literal.rs b/crates/krishiv-connect/src/translate/literal.rs index 891db16a9..c27133534 100644 --- a/crates/krishiv-connect/src/translate/literal.rs +++ b/crates/krishiv-connect/src/translate/literal.rs @@ -6,9 +6,9 @@ //! value implies; a value that does not fit a declared type is //! `INCOMPATIBLE_LITERAL_DATA_TYPE`. `NULL` is compatible with any type. -use krishiv_common::error::{KrishivError, NumericValueOutOfRangeWithSuggestion}; +use krishiv_common::error::{ErrorClass, KrishivError, NumericValueOutOfRangeWithSuggestion}; use krishiv_plan::{DayTimeField, Literal, YearMonthField}; -use krishiv_types::DecimalType; +use krishiv_types::{ArrayType, DataType, DecimalType, MapType}; use super::errors::{ assertion, incompatible_literal, not_set, unsupported_literal_type, unsupported_variant, @@ -30,6 +30,17 @@ pub(super) const YEAR_MONTH_DEFAULT: (i32, i32) = (0, 1); /// `CONNECT_INVALID_PLAN.*` for a literal with no value or a value that is /// out of range or does not fit its declared type. pub(crate) fn literal(lit: &expression::Literal) -> Result { + literal_at(lit, None, 0) +} + +/// A literal inside a nested one: its type is the one the enclosing type +/// gives it (`expected`), since Spark 4.1+ clients carry the whole type once, in the root's +/// `data_type`, and send no type on the members. `depth` counts nested literals. +fn literal_at( + lit: &expression::Literal, + expected: Option<&DataType>, + depth: usize, +) -> Result { let Some(value) = &lit.literal_type else { return Err(unsupported_literal_type("LITERALTYPE_NOT_SET(0)")); }; @@ -43,14 +54,20 @@ pub(crate) fn literal(lit: &expression::Literal) -> Result None, }; let field = value.variant().field; - if let Some(kind) = declared_kind + // A specialized array is the array kind's compact form. + let kind_field = declared_kind.map(|k| k.variant().field); + if let Some(kind_field) = kind_field && field != "null" - && kind.variant().field != field + && kind_field != field + && !(field == "specialized_array" && kind_field == "array") { - return Err(incompatible_literal(kind.variant().field, field)); + return Err(incompatible_literal(kind_field, field)); } Ok(match value { - LiteralType::Null(own) => null(declared.unwrap_or(own))?, + LiteralType::Null(own) => match (declared, expected) { + (None, Some(t)) => Literal::Null(Some(t.clone()).filter(|t| *t != DataType::Null)), + _ => null(declared.unwrap_or(own))?, + }, LiteralType::Boolean(v) => Literal::Boolean(*v), LiteralType::Byte(v) => Literal::Byte( i8::try_from(*v).map_err(|_| assertion("a byte literal is outside -128..=127"))?, @@ -63,7 +80,14 @@ pub(crate) fn literal(lit: &expression::Literal) -> Result Literal::Long(*v), LiteralType::Float(v) => Literal::Float(*v), LiteralType::Double(v) => Literal::Double(*v), - LiteralType::Decimal(d) => decimal(d, declared_kind)?, + LiteralType::Decimal(d) => { + let declared = match (declared_kind, expected) { + (Some(Kind::Decimal(t)), _) => Some(decimal_type(t.precision, t.scale)?), + (None, Some(DataType::Decimal(t))) => Some(*t), + _ => None, + }; + decimal(d, declared)? + } LiteralType::String(v) => { if let Some(Kind::String(s)) = declared_kind { // A string literal declared with a collation is refused until Connect types are gated @@ -81,13 +105,17 @@ pub(crate) fn literal(lit: &expression::Literal) -> Result Literal::Timestamp(*v), LiteralType::TimestampNtz(v) => Literal::TimestampNtz(*v), LiteralType::YearMonthInterval(months) => { - let (start, end) = match declared_kind { - Some(Kind::YearMonthInterval(t)) => { - fields(t.start_field, t.end_field, YEAR_MONTH_DEFAULT) + let (start, end) = match (declared_kind, expected) { + (Some(Kind::YearMonthInterval(t)), _) => { + let (s, e) = fields(t.start_field, t.end_field, YEAR_MONTH_DEFAULT); + (year_month_field(s)?, year_month_field(e)?) } - _ => YEAR_MONTH_DEFAULT, + (None, Some(DataType::YearMonthInterval { start, end })) => (*start, *end), + _ => ( + year_month_field(YEAR_MONTH_DEFAULT.0)?, + year_month_field(YEAR_MONTH_DEFAULT.1)?, + ), }; - let (start, end) = (year_month_field(start)?, year_month_field(end)?); if start > end { return Err(assertion("an interval literal starts after it ends")); } @@ -98,13 +126,17 @@ pub(crate) fn literal(lit: &expression::Literal) -> Result { - let (start, end) = match declared_kind { - Some(Kind::DayTimeInterval(t)) => { - fields(t.start_field, t.end_field, DAY_TIME_DEFAULT) + let (start, end) = match (declared_kind, expected) { + (Some(Kind::DayTimeInterval(t)), _) => { + let (s, e) = fields(t.start_field, t.end_field, DAY_TIME_DEFAULT); + (day_time_field(s)?, day_time_field(e)?) } - _ => DAY_TIME_DEFAULT, + (None, Some(DataType::DayTimeInterval { start, end })) => (*start, *end), + _ => ( + day_time_field(DAY_TIME_DEFAULT.0)?, + day_time_field(DAY_TIME_DEFAULT.1)?, + ), }; - let (start, end) = (day_time_field(start)?, day_time_field(end)?); if start > end { return Err(assertion("an interval literal starts after it ends")); } @@ -119,11 +151,23 @@ pub(crate) fn literal(lit: &expression::Literal) -> Result { + let root = root_type(lit, expected)?; + array(a, root.as_ref(), depth)? + } + LiteralType::Map(m) => { + let root = root_type(lit, expected)?; + map(m, root.as_ref(), depth)? + } + LiteralType::Struct(s) => { + let root = root_type(lit, expected)?; + struct_literal(s, root.as_ref(), depth)? + } + LiteralType::SpecializedArray(a) => { + let root = root_type(lit, expected)?; + specialized(a, root.as_ref())? + } + LiteralType::Time(_) | LiteralType::TimestampNtzNanos(_) | LiteralType::TimestampLtzNanos(_) => { return Err(unsupported_variant( @@ -134,6 +178,195 @@ pub(crate) fn literal(lit: &expression::Literal) -> Result, +) -> Result, KrishivError> { + match &lit.data_type { + Some(t) => Ok(Some(data_type(t)?)), + None => Ok(expected.cloned()), + } +} + +fn too_deep(depth: usize) -> Result<(), KrishivError> { + if depth >= DataType::MAX_NESTING_DEPTH { + return Err(krishiv_common::error::FailedToParseTooComplex {}.into()); + } + Ok(()) +} + +/// A type from a message's own type field when the root carries none: an unset one is the +/// message's missing type (golden-pending: Spark infers it from the elements). +fn own_type(t: Option<&ProtoType>, message: &str) -> Result { + match t { + Some(t) => data_type(t), + None => Err(assertion(format!( + "a {message} literal needs a type (the literal's data_type or its own)" + ))), + } +} + +/// `Literal.Array`: the element type from the declared array type, else the message's +/// `element_type` (`containsNull` true, as `ArrayType(t)`). +#[expect( + deprecated, + reason = "clients before Spark 4.1 type a nested literal by the message's own type fields" +)] +fn array( + a: &expression::literal::Array, + declared: Option<&DataType>, + depth: usize, +) -> Result { + too_deep(depth)?; + let ty = match declared { + Some(DataType::Array(t)) => (**t).clone(), + Some(other) => { + return Err(incompatible_literal(&other.type_name(), "array")); + } + None => ArrayType::new(own_type(a.element_type.as_ref(), "array")?, true), + }; + let elements = a + .elements + .iter() + .map(|e| literal_at(e, Some(ty.element_type()), depth + 1)) + .collect::, _>>()?; + client_input(Literal::array(ty, elements)) +} + +/// `Literal.Map`: keys and values are parallel lists of one length. +#[expect( + deprecated, + reason = "clients before Spark 4.1 type a nested literal by the message's own type fields" +)] +fn map( + m: &expression::literal::Map, + declared: Option<&DataType>, + depth: usize, +) -> Result { + too_deep(depth)?; + let ty = match declared { + Some(DataType::Map(t)) => (**t).clone(), + Some(other) => { + return Err(incompatible_literal(&other.type_name(), "map")); + } + None => MapType::new( + own_type(m.key_type.as_ref(), "map")?, + own_type(m.value_type.as_ref(), "map")?, + true, + ), + }; + if m.keys.len() != m.values.len() { + return Err(assertion( + "a map literal has a different number of keys and values", + )); + } + let entries = m + .keys + .iter() + .zip(&m.values) + .map(|(k, v)| { + Ok(( + literal_at(k, Some(ty.key_type()), depth + 1)?, + literal_at(v, Some(ty.value_type()), depth + 1)?, + )) + }) + .collect::, KrishivError>>()?; + client_input(Literal::key_value_map(ty, entries)) +} + +/// `Literal.Struct`: one element per field of the struct type. +#[expect( + deprecated, + reason = "clients before Spark 4.1 type a nested literal by the message's own type fields" +)] +fn struct_literal( + s: &expression::literal::Struct, + declared: Option<&DataType>, + depth: usize, +) -> Result { + too_deep(depth)?; + let ty = match declared { + Some(DataType::Struct(t)) => t.clone(), + Some(other) => { + return Err(incompatible_literal(&other.type_name(), "struct")); + } + None => match own_type(s.struct_type.as_ref(), "struct")? { + DataType::Struct(t) => t, + other => return Err(incompatible_literal(&other.type_name(), "struct")), + }, + }; + if ty.fields().len() != s.elements.len() { + return Err(assertion( + "a struct literal has a different number of elements and fields", + )); + } + let fields = ty + .fields() + .iter() + .zip(&s.elements) + .map(|(f, e)| literal_at(e, Some(f.data_type()), depth + 1)) + .collect::, _>>()?; + client_input(Literal::struct_value(ty, fields)) +} + +/// `Literal.SpecializedArray`: a primitive array, as an array literal of non-null elements. +/// The root type, when there is one, must be the array of that primitive. +fn specialized( + a: &expression::literal::SpecializedArray, + declared: Option<&DataType>, +) -> Result { + use expression::literal::specialized_array::ValueType as V; + let Some(values) = &a.value_type else { + return Err(not_set( + "spark.connect.Expression.Literal.SpecializedArray", + "VALUETYPE_NOT_SET", + )); + }; + let (element, elements): (DataType, Vec) = match values { + V::Bools(v) => ( + DataType::Boolean, + v.values.iter().map(|x| Literal::Boolean(*x)).collect(), + ), + V::Ints(v) => ( + DataType::Integer, + v.values.iter().map(|x| Literal::Integer(*x)).collect(), + ), + V::Longs(v) => ( + DataType::Long, + v.values.iter().map(|x| Literal::Long(*x)).collect(), + ), + V::Floats(v) => ( + DataType::Float, + v.values.iter().map(|x| Literal::Float(*x)).collect(), + ), + V::Doubles(v) => ( + DataType::Double, + v.values.iter().map(|x| Literal::Double(*x)).collect(), + ), + V::Strings(v) => ( + DataType::String, + v.values + .iter() + .map(|x| Literal::String(x.clone())) + .collect(), + ), + }; + let ty = match declared { + // The declared type may allow nulls the values do not have. + Some(DataType::Array(t)) => (**t).clone(), + Some(other) => { + return Err(incompatible_literal( + &other.type_name(), + "specialized_array", + )); + } + None => ArrayType::new(element, false), + }; + client_input(Literal::array(ty, elements)) +} + /// A null of type `of`; the `NULL` type is Spark's `void`. fn null(of: &ProtoType) -> Result { match &of.kind { @@ -180,12 +413,12 @@ pub(super) fn day_time_field(n: i32) -> Result { /// precision and scale the message gives (Spark's `getProtoDataTypeDefault`). /// The unscaled value is `value` at that scale, rounded half up if the /// declared scale is smaller than the value's own. -fn decimal(d: &Decimal, declared: Option<&Kind>) -> Result { +fn decimal(d: &Decimal, declared: Option) -> Result { let parsed = parse_decimal(&d.value).ok_or_else(|| assertion("a decimal literal is not a number"))?; let data_type = match declared { - Some(Kind::Decimal(t)) => decimal_type(t.precision, t.scale)?, - _ => { + Some(t) => t, + None => { let widest = |own: i64, given: Option| given.map_or(own, |g| own.max(i64::from(g))); let scale = widest(parsed.scale, d.scale); @@ -341,5 +574,18 @@ fn parse_decimal(text: &str) -> Option { }) } +/// A nested literal the plan crate refuses for a member that does not fit the declared type (or +/// a null where nulls are forbidden): client input, so `CONNECT_INVALID_PLAN.ASSERTION_FAILURE` +/// rather than the constructor's `INTERNAL_ERROR`. Other errors (too deep) pass through. +fn client_input(r: Result) -> Result { + r.map_err(|e| { + if e.error_class() == ErrorClass::InternalError { + assertion(e.to_string()) + } else { + e + } + }) +} + #[cfg(test)] mod tests; diff --git a/crates/krishiv-connect/src/translate/literal/tests.rs b/crates/krishiv-connect/src/translate/literal/tests.rs index 408384ab4..4027bdb0e 100644 --- a/crates/krishiv-connect/src/translate/literal/tests.rs +++ b/crates/krishiv-connect/src/translate/literal/tests.rs @@ -509,30 +509,238 @@ fn a_calendar_interval_literal_keeps_its_parts() { } #[test] -fn literals_the_ir_has_no_value_for_are_unsupported() { - for (value, name) in [ - (LiteralType::Array(pl::Array::default()), "array"), - (LiteralType::Map(pl::Map::default()), "map"), - (LiteralType::Struct(pl::Struct::default()), "struct"), - ( - LiteralType::SpecializedArray(pl::SpecializedArray::default()), - "specialized_array", - ), - (LiteralType::Time(pl::Time::default()), "time"), - ] { - let e = literal(&literal_of(value)).expect_err(name); - assert_eq!( - e.error_class(), - ErrorClass::ConnectInvalidPlanInvalidOneOfFieldNotSupported, - "{name}" - ); - assert!( - e.message().contains("spark.connect.Expression.Literal") - && e.message().contains(&name.to_ascii_uppercase()), - "{name}: {}", - e.message() - ); +fn a_time_literal_is_unsupported() { + let e = literal(&literal_of(LiteralType::Time(pl::Time::default()))).expect_err("time"); + assert_eq!( + e.error_class(), + ErrorClass::ConnectInvalidPlanInvalidOneOfFieldNotSupported + ); + assert!( + e.message().contains("spark.connect.Expression.Literal") && e.message().contains("TIME"), + "{}", + e.message() + ); +} + +fn int_lit(v: i32) -> expression::Literal { + literal_of(LiteralType::Integer(v)) +} + +fn null_lit() -> expression::Literal { + literal_of(LiteralType::Null(proto_type(Kind::Integer( + dt::Integer::default(), + )))) +} + +fn array_type(element: Kind, contains_null: bool) -> DataTypeProto { + proto_type(Kind::Array(Box::new(dt::Array { + element_type: Some(Box::new(proto_type(element))), + contains_null, + ..dt::Array::default() + }))) +} + +#[test] +fn an_array_literal_takes_its_type_from_the_root_data_type() { + // A Spark 4.1+ client: the root carries the type, the elements none. + let lit = expression::Literal { + literal_type: Some(LiteralType::Array(pl::Array { + elements: vec![int_lit(1), null_lit(), int_lit(3)], + ..pl::Array::default() + })), + data_type: Some(array_type(Kind::Integer(dt::Integer::default()), true)), + }; + let got = literal(&lit).unwrap(); + let Literal::Array { + data_type, + elements, + } = &got + else { + panic!("not an array: {got:?}"); + }; + assert_eq!(data_type.element_type(), &DataType::Integer); + assert!(data_type.contains_null()); + assert_eq!( + elements, + &[ + Literal::Integer(1), + Literal::Null(Some(DataType::Integer)), + Literal::Integer(3) + ] + ); + // A null where the type forbids one is refused. + let strict = expression::Literal { + data_type: Some(array_type(Kind::Integer(dt::Integer::default()), false)), + ..lit + }; + assert!(literal(&strict).is_err()); +} + +#[test] +#[expect(deprecated, reason = "the pre-4.1 element_type field")] +fn an_array_literal_of_an_older_client_names_its_element_type() { + let lit = literal_of(LiteralType::Array(pl::Array { + element_type: Some(proto_type(Kind::Integer(dt::Integer::default()))), + elements: vec![int_lit(7)], + })); + let got = literal(&lit).unwrap(); + assert_eq!(got.data_type().unwrap().simple_string(), "array"); + // Neither a root type nor an element type: Spark would infer one; Krishiv refuses. + let none = literal_of(LiteralType::Array(pl::Array { + elements: vec![int_lit(7)], + ..pl::Array::default() + })); + assert_eq!( + literal(&none).unwrap_err().error_class(), + ErrorClass::ConnectInvalidPlanAssertionFailure + ); +} + +#[test] +fn an_empty_array_literal_is_an_array_of_no_elements() { + let lit = expression::Literal { + literal_type: Some(LiteralType::Array(pl::Array::default())), + data_type: Some(array_type(Kind::String(dt::String::default()), true)), + }; + let got = literal(&lit).unwrap(); + assert_eq!(got.data_type().unwrap().simple_string(), "array"); +} + +#[test] +fn a_map_literal_pairs_keys_with_values() { + let map_type = proto_type(Kind::Map(Box::new(dt::Map { + key_type: Some(Box::new(proto_type(Kind::String(dt::String::default())))), + value_type: Some(Box::new(proto_type(Kind::Integer(dt::Integer::default())))), + value_contains_null: true, + ..dt::Map::default() + }))); + let lit = expression::Literal { + literal_type: Some(LiteralType::Map(pl::Map { + keys: vec![literal_of(LiteralType::String("a".into()))], + values: vec![int_lit(1)], + ..pl::Map::default() + })), + data_type: Some(map_type.clone()), + }; + let got = literal(&lit).unwrap(); + assert_eq!(got.data_type().unwrap().simple_string(), "map"); + // A different number of keys and values is an invalid plan. + let bad = expression::Literal { + literal_type: Some(LiteralType::Map(pl::Map { + keys: vec![literal_of(LiteralType::String("a".into()))], + ..pl::Map::default() + })), + data_type: Some(map_type), + }; + assert_eq!( + literal(&bad).unwrap_err().error_class(), + ErrorClass::ConnectInvalidPlanAssertionFailure + ); +} + +#[test] +fn a_struct_literal_has_one_element_per_field_and_nests() { + let field = |name: &str, kind: Kind| dt::StructField { + name: name.to_owned(), + data_type: Some(proto_type(kind)), + nullable: true, + metadata: None, + }; + let inner = Kind::Array(Box::new(dt::Array { + element_type: Some(Box::new(proto_type(Kind::Integer(dt::Integer::default())))), + contains_null: false, + ..dt::Array::default() + })); + let struct_type = proto_type(Kind::Struct(dt::Struct { + fields: vec![ + field("a", Kind::Integer(dt::Integer::default())), + field("b", inner), + ], + ..dt::Struct::default() + })); + let lit = expression::Literal { + literal_type: Some(LiteralType::Struct(pl::Struct { + elements: vec![ + int_lit(1), + literal_of(LiteralType::Array(pl::Array { + elements: vec![int_lit(2)], + ..pl::Array::default() + })), + ], + ..pl::Struct::default() + })), + data_type: Some(struct_type.clone()), + }; + let got = literal(&lit).unwrap(); + assert_eq!( + got.data_type().unwrap().simple_string(), + "struct>" + ); + // The wrong count of elements. + let short = expression::Literal { + literal_type: Some(LiteralType::Struct(pl::Struct::default())), + data_type: Some(struct_type), + }; + assert_eq!( + literal(&short).unwrap_err().error_class(), + ErrorClass::ConnectInvalidPlanAssertionFailure + ); +} + +#[test] +fn a_specialized_array_is_an_array_literal() { + let lit = literal_of(LiteralType::SpecializedArray(pl::SpecializedArray { + value_type: Some(pl::specialized_array::ValueType::Longs( + crate::proto::v4_3::Longs { values: vec![1, 2] }, + )), + })); + let got = literal(&lit).unwrap(); + assert_eq!(got.data_type().unwrap().simple_string(), "array"); + // With the root type of the same array it is accepted too, and a different kind is refused. + let typed_root = expression::Literal { + data_type: Some(array_type(Kind::Long(dt::Long::default()), true)), + ..lit.clone() + }; + assert!(literal(&typed_root).is_ok()); + let wrong = expression::Literal { + data_type: Some(proto_type(Kind::Integer(dt::Integer::default()))), + ..lit + }; + assert_eq!( + literal(&wrong).unwrap_err().error_class(), + ErrorClass::ConnectInvalidPlanIncompatibleLiteralDataType + ); +} + +#[test] +fn a_nested_literal_deeper_than_the_type_limit_is_too_complex() { + // array>> past MAX_NESTING_DEPTH levels, each level a literal. + let n = DataType::MAX_NESTING_DEPTH + 2; + let mut ty = proto_type(Kind::Integer(dt::Integer::default())); + let mut value = int_lit(1); + for _ in 0..n { + ty = proto_type(Kind::Array(Box::new(dt::Array { + element_type: Some(Box::new(ty)), + contains_null: false, + ..dt::Array::default() + }))); + value = literal_of(LiteralType::Array(pl::Array { + elements: vec![value], + ..pl::Array::default() + })); } + value.data_type = Some(ty); + let e = std::thread::Builder::new() + .stack_size(64 << 20) + .spawn(move || { + // The deep message drops here, on the big stack. + literal(&value) + }) + .unwrap() + .join() + .unwrap() + .unwrap_err(); + assert_eq!(e.error_class(), ErrorClass::FailedToParseTooComplex); } #[test] @@ -573,3 +781,34 @@ fn a_decimal_literal_of_any_length_is_read_in_place() { ErrorClass::DecimalPrecisionExceedsMaxPrecision ); } + +#[test] +#[expect(deprecated, reason = "the pre-4.1 element_type field")] +fn a_member_that_misses_the_declared_type_is_an_invalid_plan_not_internal() { + let assertion = ErrorClass::ConnectInvalidPlanAssertionFailure; + // A pre-4.1 client declares `int` elements and sends a long. + let wrong_member = literal_of(LiteralType::Array(pl::Array { + element_type: Some(proto_type(Kind::Integer(dt::Integer::default()))), + elements: vec![literal_of(LiteralType::Long(1))], + })); + assert_eq!(literal(&wrong_member).unwrap_err().error_class(), assertion); + // A null member under `array`. + let null_member = expression::Literal { + literal_type: Some(LiteralType::Array(pl::Array { + elements: vec![int_lit(1), null_lit()], + ..pl::Array::default() + })), + data_type: Some(array_type(Kind::Integer(dt::Integer::default()), false)), + }; + assert_eq!(literal(&null_member).unwrap_err().error_class(), assertion); + // A specialized `longs` under a root `array`. + let specialized = expression::Literal { + literal_type: Some(LiteralType::SpecializedArray(pl::SpecializedArray { + value_type: Some(pl::specialized_array::ValueType::Longs( + crate::proto::v4_3::Longs { values: vec![1] }, + )), + })), + data_type: Some(array_type(Kind::Integer(dt::Integer::default()), true)), + }; + assert_eq!(literal(&specialized).unwrap_err().error_class(), assertion); +} diff --git a/crates/krishiv-connect/src/translate/local.rs b/crates/krishiv-connect/src/translate/local.rs new file mode 100644 index 000000000..a4ca28a54 --- /dev/null +++ b/crates/krishiv-connect/src/translate/local.rs @@ -0,0 +1,162 @@ +// SPDX-License-Identifier: Apache-2.0 +//! `LocalRelation` and `CachedLocalRelation` (`createDataFrame`): the host's local relations. +//! +//! The translator is synchronous and building a relation is not (decoding the Arrow IPC data and +//! applying a user schema run on the session's executor), so the host runs an async pre-pass +//! (design P3-T012a part 2b, re-review nit N4): it translates the plan once with +//! [`LocalRelations::collecting`], which records a [`LocalRequest`] per local relation and +//! answers with placeholders; builds the relations (`Session::local_relation`); and translates +//! again with [`LocalRelations::resolved`], which hands the translator the n-th relation for the +//! n-th request. A plan without local relations needs the first translation only. Translation +//! visits the relations in a fixed order, so the n-th request of both passes is the same relation. + +use std::sync::{Arc, Mutex, PoisonError}; + +use krishiv_common::error::{ + ConnectInvalidPlanSchemaRequiredForLocalRelation, KrishivError, KrishivUnsupportedFeature, +}; +use krishiv_plan::LocalRelation; +use krishiv_types::StructType; + +use super::schema::parse_local_schema; +use crate::proto::v4_3 as pb; + +/// A local relation the host builds before the second translation. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum LocalRequest { + /// `LocalRelation`: Arrow IPC data with its own schema and a user schema, either of which may + /// be absent (not both). + Inline { + /// The Arrow IPC stream. + data: Option>, + /// The parsed `schema` (DDL or JSON; a type that is not a struct is `struct`). + schema: Option, + }, + /// `CachedLocalRelation`: the artifact `cache/` holding a serialized `LocalRelation`. + Cached { + /// The hash the client names. + hash: String, + }, +} + +/// Where the translator gets its local relations from. +#[derive(Clone, Default)] +pub struct LocalRelations(Mode); + +#[derive(Clone, Default)] +enum Mode { + /// No host: a local relation is refused. + #[default] + None, + /// The first pass: record the requests, answer with placeholders. + Collect(Arc>>), + /// The second pass: the relation of the n-th request. + Resolved(Arc<[LocalRelation]>), +} + +impl LocalRelations { + /// The first pass: records what the plan needs. + #[must_use] + pub fn collecting() -> Self { + Self(Mode::Collect(Arc::default())) + } + + /// The second pass: `relations[n]` answers the n-th request of the first pass. + #[must_use] + pub fn resolved(relations: Vec) -> Self { + Self(Mode::Resolved(relations.into())) + } + + /// Takes the requests recorded so far, in order (empty unless [`Self::collecting`]), leaving + /// none behind: the inline data is moved out, never copied. + #[must_use] + pub fn take_requests(&self) -> Vec { + match &self.0 { + Mode::Collect(r) => { + std::mem::take(&mut *r.lock().unwrap_or_else(PoisonError::into_inner)) + } + _ => Vec::new(), + } + } + + /// The relation for the request `ordinal`: recorded (and a placeholder returned) in the first + /// pass, looked up in the second. + pub(super) fn get( + &self, + ordinal: usize, + request: impl FnOnce() -> Result, + ) -> Result { + match &self.0 { + Mode::None => Err(KrishivUnsupportedFeature { + feature: "local relations without a host that builds them \ + (TranslateConf::local_relations)" + .to_owned(), + } + .into()), + Mode::Collect(requests) => { + let request = request()?; + requests + .lock() + .unwrap_or_else(PoisonError::into_inner) + .push(request); + // Never run: the first pass's plan is dropped when there were requests. + Ok(LocalRelation::from_decoded( + Arc::new(arrow_schema::Schema::empty()), + [0; 32], + 0, + 0, + 0, + )) + } + Mode::Resolved(relations) => relations.get(ordinal).cloned().ok_or_else(|| { + krishiv_common::error::InternalError { + message: format!("no local relation was built for request {ordinal}"), + } + .into() + }), + } + } +} + +impl std::fmt::Debug for LocalRelations { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match &self.0 { + Mode::None => f.write_str("LocalRelations::None"), + Mode::Collect(_) => f.write_str("LocalRelations::Collect"), + Mode::Resolved(r) => write!(f, "LocalRelations::Resolved({})", r.len()), + } + } +} + +impl PartialEq for LocalRelations { + fn eq(&self, other: &Self) -> bool { + match (&self.0, &other.0) { + (Mode::None, Mode::None) => true, + (Mode::Collect(a), Mode::Collect(b)) => Arc::ptr_eq(a, b), + (Mode::Resolved(a), Mode::Resolved(b)) => a == b, + _ => false, + } + } +} + +impl Eq for LocalRelations {} + +/// The request of a `LocalRelation`, as Spark's `transformLocalRelation` reads it: a schema, +/// when set, is parsed (DDL, else JSON); with neither data nor schema the relation is invalid. +/// +/// # Errors +/// `CONNECT_INVALID_PLAN.SCHEMA_REQUIRED_FOR_LOCAL_RELATION` for neither, the parse error of a +/// schema that is neither DDL nor JSON. +pub(super) fn inline_request(l: &pb::LocalRelation) -> Result { + let schema = l.schema.as_deref().map(parse_local_schema).transpose()?; + if l.data.is_none() && schema.is_none() { + return Err(ConnectInvalidPlanSchemaRequiredForLocalRelation {}.into()); + } + Ok(LocalRequest::Inline { + data: l.data.clone(), + schema, + }) +} + +#[cfg(test)] +mod tests; diff --git a/crates/krishiv-connect/src/translate/local/tests.rs b/crates/krishiv-connect/src/translate/local/tests.rs new file mode 100644 index 000000000..2ad9a80f6 --- /dev/null +++ b/crates/krishiv-connect/src/translate/local/tests.rs @@ -0,0 +1,144 @@ +// SPDX-License-Identifier: Apache-2.0 +//! Tests of the local relation arms: the host protocol (requests, ordinals) and the refusals. + +use std::sync::Arc; + +use krishiv_common::error::ErrorClass; +use krishiv_plan::{LocalRelation, PlanNode, ReadSource}; + +use super::super::testing::{boxed, fails, relation}; +use super::super::{TranslateConf, translate_relation}; +use super::*; +use crate::proto::v4_3::{Join, Relation, join::JoinType, relation::RelType}; + +fn local(data: Option<&[u8]>, schema: Option<&str>) -> Relation { + relation(RelType::LocalRelation(pb::LocalRelation { + data: data.map(<[u8]>::to_vec), + schema: schema.map(str::to_owned), + })) +} + +fn cached(hash: &str) -> Relation { + relation(RelType::CachedLocalRelation(pb::CachedLocalRelation { + hash: hash.to_owned(), + })) +} + +fn conf_with(local_relations: LocalRelations) -> TranslateConf { + TranslateConf { + local_relations, + ..TranslateConf::default() + } +} + +fn built(rows: u64) -> LocalRelation { + LocalRelation::from_decoded(Arc::new(arrow_schema::Schema::empty()), [7; 32], 1, rows, 0) +} + +#[test] +fn the_first_pass_records_each_request_in_order() { + let locals = LocalRelations::collecting(); + let join = relation(RelType::Join(Box::new(Join { + left: boxed(local(Some(b"abc"), Some("a INT"))), + right: boxed(cached("h1")), + join_type: JoinType::Cross as i32, + ..Join::default() + }))); + translate_relation(&join, &conf_with(locals.clone())).expect("translates"); + let requests = locals.take_requests(); + assert_eq!(requests.len(), 2); + let LocalRequest::Inline { data, schema } = &requests[0] else { + panic!("{:?}", requests[0]) + }; + assert_eq!(data.as_deref(), Some(&b"abc"[..])); + assert_eq!(schema.as_ref().map(|s| s.fields().len()), Some(1)); + assert_eq!(requests[1], LocalRequest::Cached { hash: "h1".into() }); +} + +#[test] +fn the_second_pass_hands_out_the_relation_of_each_request() { + let join = relation(RelType::Join(Box::new(Join { + left: boxed(local(Some(b"abc"), None)), + right: boxed(local(None, Some("b LONG"))), + join_type: JoinType::Cross as i32, + ..Join::default() + }))); + let conf = conf_with(LocalRelations::resolved(vec![built(1), built(2)])); + let plan = translate_relation(&join, &conf).expect("translates"); + let PlanNode::Join { left, right, .. } = &plan.root().node else { + panic!("not a join") + }; + let rows = |p: &Arc| match &p.node { + PlanNode::Read(ReadSource::Local(l)) => l.rows(), + other => panic!("{other:?}"), + }; + assert_eq!((rows(left), rows(right)), (1, 2)); + // A host that built too few is an internal error, never a silent empty relation. + let few = conf_with(LocalRelations::resolved(vec![built(1)])); + assert_eq!( + fails(translate_relation(&join, &few)).error_class(), + ErrorClass::InternalError + ); +} + +#[test] +fn no_data_and_no_schema_is_an_invalid_relation() { + let e = fails(translate_relation( + &local(None, None), + &conf_with(LocalRelations::collecting()), + )); + assert_eq!( + e.error_class().name(), + "CONNECT_INVALID_PLAN.SCHEMA_REQUIRED_FOR_LOCAL_RELATION" + ); + // A schema alone is an empty relation of that schema. + let locals = LocalRelations::collecting(); + translate_relation(&local(None, Some("a INT")), &conf_with(locals.clone())).expect("ok"); + assert_eq!(locals.take_requests().len(), 1); +} + +#[test] +fn a_bad_schema_text_fails_at_translation_and_is_not_echoed() { + let e = fails(translate_relation( + &local(Some(b"x"), Some("CANARY nonsense nonsense")), + &conf_with(LocalRelations::collecting()), + )); + assert!(!e.message().contains("CANARY"), "{}", e.message()); +} + +#[test] +fn a_host_that_builds_none_refuses_local_relations() { + for r in [local(Some(b"x"), None), cached("h")] { + let e = fails(translate_relation(&r, &TranslateConf::default())); + assert_eq!(e.error_class().name(), "KRISHIV_UNSUPPORTED_FEATURE"); + } +} + +#[test] +fn the_chunked_cached_relation_stays_refused() { + let r = relation(RelType::ChunkedCachedLocalRelation( + pb::ChunkedCachedLocalRelation::default(), + )); + let e = fails(translate_relation( + &r, + &conf_with(LocalRelations::collecting()), + )); + assert_eq!( + e.error_class(), + ErrorClass::ConnectInvalidPlanInvalidOneOfFieldNotSupported + ); +} + +#[test] +fn local_relations_compare_by_mode() { + assert_eq!(LocalRelations::default(), LocalRelations::default()); + let c = LocalRelations::collecting(); + assert_eq!(c, c.clone()); + assert_ne!(c, LocalRelations::collecting()); + assert_eq!( + LocalRelations::resolved(vec![built(1)]), + LocalRelations::resolved(vec![built(1)]) + ); + assert_ne!(c, LocalRelations::default()); + assert!(format!("{c:?}").contains("Collect")); +} diff --git a/crates/krishiv-connect/src/translate/mod.rs b/crates/krishiv-connect/src/translate/mod.rs index 6e793b708..38da9c898 100644 --- a/crates/krishiv-connect/src/translate/mod.rs +++ b/crates/krishiv-connect/src/translate/mod.rs @@ -76,10 +76,15 @@ mod errors; mod expression; mod literal; +mod local; mod names; mod relation; +mod schema; mod types; +pub use local::{LocalRelations, LocalRequest}; +pub use schema::{parse_datatype_string, parse_local_schema, parse_table_schema}; + use krishiv_common::error::{FailedToParseTooComplex, KrishivError}; use krishiv_plan::{Expr, MAX_DEPTH, Plan}; use std::sync::Arc; @@ -97,6 +102,11 @@ type Result = std::result::Result; pub struct TranslateConf { /// The format of a read that names none: `spark.sql.sources.default`. pub default_data_source: String, + /// The session settings that change how SQL text parses; an + /// `ExpressionString` (`expr()`, `selectExpr`) is parsed with them. + pub parser: krishiv_sql::ParserOptions, + /// The local relations of the plan (`createDataFrame`), built by the host. + pub local_relations: LocalRelations, } impl Default for TranslateConf { @@ -104,6 +114,8 @@ impl Default for TranslateConf { fn default() -> Self { Self { default_data_source: "parquet".to_owned(), + parser: krishiv_sql::ParserOptions::default(), + local_relations: LocalRelations::default(), } } } @@ -132,18 +144,23 @@ pub fn expression_from_bytes(bytes: &[u8], conf: &TranslateConf) -> Result Result { match &plan.op_type { Some(OpType::Root(relation)) => translate_relation(relation, conf), - Some(OpType::Command(command)) => Err(match &command.command_type { - Some(arm) => unsupported_variant("spark.connect.Command", arm.variant()), - None => not_set("spark.connect.Command", "COMMAND_TYPE_NOT_SET"), - }), + Some(OpType::Command(command)) => match &command.command_type { + Some(pb::command::CommandType::SqlCommand(c)) => { + translate_relation(&sql_command_relation(c)?, conf) + } + Some(arm) => Err(unsupported_variant("spark.connect.Command", arm.variant())), + None => Err(not_set("spark.connect.Command", "COMMAND_TYPE_NOT_SET")), + }, Some(OpType::CompressedOperation(_)) => Err(unsupported_field( "Plan", "compressed_operation", @@ -153,6 +170,66 @@ pub fn translate_plan(plan: &pb::Plan, conf: &TranslateConf) -> Result { } } +/// The `SQL` relation of a `SqlCommand`: what `spark.sql()`'s client re-sends when the statement +/// was a query. When `input` is set (every `spark.sql()` of a PySpark 4.x client sends it, the +/// other fields being deprecated) it is the relation, unchanged and with its `common` (Spark's +/// `handleSqlCommand` answers `command.getInput`); otherwise the deprecated `sql`, `args` and +/// `pos_args` fields are the relation's `query`, `args` and `pos_args`, the named and positional +/// expressions carrying over as they are. +/// +/// The library form of a `sql_command` exists for the coverage tie and for the server, which +/// runs the command; `AnalyzePlan` of a command plan stays refused by the server. +/// +/// # Errors +/// `KRISHIV_UNSUPPORTED_FEATURE` for an `input` of `WithRelations` (SQL over DataFrame +/// arguments, Connect group 4); `CONNECT_INVALID_PLAN.SQL_COMMAND_EXPECTS_SQL_OR_WITH_RELATIONS` +/// for any other `input` relation. +pub fn sql_command_relation(command: &pb::SqlCommand) -> Result { + if let Some(input) = &command.input { + return match &input.rel_type { + Some(pb::relation::RelType::Sql(_)) => Ok(input.clone()), + Some(pb::relation::RelType::WithRelations(_)) => Err(unsupported_field( + "SqlCommand", + "input", + "SQL over DataFrame arguments", + )), + other => Err( + krishiv_common::error::ConnectInvalidPlanSqlCommandExpectsSqlOrWithRelations { + other: other.as_ref().map_or_else( + || "REL_TYPE_NOT_SET".to_owned(), + |r| r.variant().field.to_ascii_uppercase(), + ), + } + .into(), + ), + }; + } + #[expect( + deprecated, + reason = "Spark still accepts the deprecated fields of older clients" + )] + let sql = pb::Sql { + query: command.sql.clone(), + args: command.args.clone(), + pos_args: command.pos_args.clone(), + named_arguments: command.named_arguments.clone(), + pos_arguments: command.pos_arguments.clone(), + }; + Ok(pb::Relation { + common: None, + rel_type: Some(pb::relation::RelType::Sql(sql)), + }) +} + +/// Whether a `SQL` relation carries parameters (`args`, `pos_args`, named or positional +/// expressions). +#[must_use] +pub fn sql_has_parameters(sql: &pb::Sql) -> bool { + #[expect(deprecated, reason = "older clients send the literal fields")] + let literals = !sql.args.is_empty() || !sql.pos_args.is_empty(); + literals || !sql.named_arguments.is_empty() || !sql.pos_arguments.is_empty() +} + /// The refusal of a command (`Plan.command`) until the server translates /// commands (P3-T012): Spark's `CONNECT_INVALID_PLAN.INVALID_ONE_OF_FIELD_NOT_SUPPORTED` /// naming the command's case (`SQL_COMMAND`, 56K00), or @@ -205,6 +282,8 @@ pub(crate) struct Translator<'c> { nodes: usize, /// Threads started so far for subtrees past [`INLINE_LEVELS`]. hops: usize, + /// Local relations met so far: the ordinal of the next one's request. + locals: usize, } /// The most threads one translation starts for deep subtrees. A plan needs @@ -221,6 +300,7 @@ impl<'c> Translator<'c> { depth: 0, nodes: 0, hops: 0, + locals: 0, } } diff --git a/crates/krishiv-connect/src/translate/relation.rs b/crates/krishiv-connect/src/translate/relation.rs index 861992ab6..d67462d7d 100644 --- a/crates/krishiv-connect/src/translate/relation.rs +++ b/crates/krishiv-connect/src/translate/relation.rs @@ -13,8 +13,8 @@ use std::sync::Arc; use krishiv_common::secret::OptionValue; use krishiv_plan::{ - Expr, ExprNode, Grouping, JoinCondition, JoinKind, Literal, LogicalPlan, Origin, PlanId, - PlanNode, ReadSource, SetOperation, SqlArgs, + ALIAS_PART_SEP, Expr, ExprNode, Grouping, JoinCondition, JoinKind, Literal, LogicalPlan, + Origin, PlanId, PlanNode, ReadSource, SetOperation, SqlArgs, }; use super::errors::{ @@ -78,23 +78,24 @@ impl Translator<'_> { offset: Arc::new(Expr::literal(Literal::Integer(o.offset))), }, RelType::Aggregate(a) => self.aggregate(a)?, - RelType::SubqueryAlias(a) => { - if !a.qualifier.is_empty() { - return Err(unsupported_field( - "SubqueryAlias", - "qualifier", - "relation aliases with a qualifier", - )); - } - PlanNode::SubqueryAlias { - input: self.input("SubqueryAlias", a.input.as_deref())?, - alias: a.alias.clone(), - columns: Vec::new(), - } - } + RelType::SubqueryAlias(a) => PlanNode::SubqueryAlias { + input: self.input("SubqueryAlias", a.input.as_deref())?, + alias: alias_name(&a.qualifier, &a.alias)?, + columns: Vec::new(), + }, RelType::WithColumns(w) => self.with_columns(w)?, RelType::Range(r) => self.range(r)?, RelType::SetOp(s) => self.set_op(s, origin_of(rel))?, + RelType::LocalRelation(l) => PlanNode::Read(ReadSource::Local( + self.local_relation(|| super::local::inline_request(l))?, + )), + RelType::CachedLocalRelation(c) => { + PlanNode::Read(ReadSource::Local(self.local_relation(|| { + Ok(super::LocalRequest::Cached { + hash: c.hash.clone(), + }) + })?)) + } other => { return Err(unsupported_variant( "spark.connect.Relation", @@ -178,6 +179,16 @@ impl Translator<'_> { ) } + /// The next local relation of the plan, from the host (see [`super::LocalRelations`]). + fn local_relation( + &mut self, + request: impl FnOnce() -> Result, + ) -> Result { + let ordinal = self.locals; + self.locals += 1; + self.conf.local_relations.get(ordinal, request) + } + /// The input relation of `message`, which must be set. fn input(&mut self, message: &str, input: Option<&Relation>) -> Result> { match input { @@ -212,13 +223,11 @@ impl Translator<'_> { options: typed(&t.options), }, ReadType::DataSource(d) => { - if d.schema.as_deref().is_some_and(|s| !s.is_empty()) { - return Err(unsupported_field( - "Read.DataSource", - "schema", - "reads with a user-specified schema", - )); - } + // Spark: a set and non-empty schema is the reader's schema. + let schema = match d.schema.as_deref() { + Some(s) if !s.is_empty() => Some(super::schema::parse_table_schema(s)?), + _ => None, + }; if !d.predicates.is_empty() { return Err(unsupported_field( "Read.DataSource", @@ -226,7 +235,7 @@ impl Translator<'_> { "reads with predicates", )); } - ReadSource::files( + let source = ReadSource::files( d.format .clone() .unwrap_or_else(|| self.conf.default_data_source.clone()), @@ -235,7 +244,11 @@ impl Translator<'_> { .map(|p| OptionValue::from_user("path", p.as_str())) .collect(), typed(&d.options), - ) + ); + match schema { + Some(schema) => source.with_schema(schema), + None => source, + } } })) } @@ -464,6 +477,24 @@ impl Translator<'_> { #[cfg(test)] mod tests; +/// The alias of a `SubqueryAlias`: `qualifier ++ [alias]` joined by [`ALIAS_PART_SEP`], as the IR +/// holds a multi-part alias (Spark's `AliasIdentifier(alias, qualifier)`). A part holding the +/// separator could not be told from two parts, so it is refused as `KRISHIV_UNSUPPORTED_FEATURE` +/// (a deviation from note section 4: 56K00 is for `oneof` arms). +fn alias_name(qualifier: &[String], alias: &str) -> Result { + let parts = || qualifier.iter().map(String::as_str).chain([alias]); + if parts().any(|p| p.contains(ALIAS_PART_SEP)) { + return Err(unsupported_field( + "SubqueryAlias", + "qualifier", + "relation alias parts holding the control character U+001E", + )); + } + Ok(parts() + .collect::>() + .join(&ALIAS_PART_SEP.to_string())) +} + /// The origin of a relation: its `plan_id`, when the client sent one. fn origin_of(rel: &Relation) -> Origin { rel.common diff --git a/crates/krishiv-connect/src/translate/relation/tests.rs b/crates/krishiv-connect/src/translate/relation/tests.rs index 1900d553e..4a80f547d 100644 --- a/crates/krishiv-connect/src/translate/relation/tests.rs +++ b/crates/krishiv-connect/src/translate/relation/tests.rs @@ -86,6 +86,7 @@ fn file_reads_and_the_default_format() { assert_eq!(show(&read_files(None, &["/a"])), "Read parquet [/a]\n"); let orc = TranslateConf { default_data_source: "orc".into(), + ..TranslateConf::default() }; let plan = translate_relation(&read_files(None, &["/a"]), &orc).unwrap(); assert_eq!(plan.to_string(), "Read orc [/a]\n"); @@ -129,21 +130,12 @@ fn file_read_options_and_paths_are_classified_once() { } #[test] -fn a_streaming_read_a_schema_and_predicates_are_unsupported() { +fn a_streaming_read_and_predicates_are_unsupported() { let mut streaming = read_files(Some("parquet"), &["/a"]); if let Some(RelType::Read(r)) = &mut streaming.rel_type { r.is_streaming = true; } assert!(message(&streaming).contains("is_streaming")); - let schema = data_source(read::DataSource { - schema: Some("a INT".into()), - ..read::DataSource::default() - }); - assert!( - message(&schema).contains("Read.DataSource.schema"), - "{}", - message(&schema) - ); // The client sends an empty schema string when none is given. let empty = data_source(read::DataSource { format: Some("parquet".into()), @@ -350,7 +342,7 @@ fn range_builds_the_range_leaf() { } #[test] -fn a_project_without_input_reads_one_row_and_a_qualified_alias_is_unsupported() { +fn a_project_without_input_reads_one_row() { let no_input = relation(RelType::Project(Box::default())); match node(&no_input) { PlanNode::Project { input, exprs } => { @@ -359,12 +351,94 @@ fn a_project_without_input_reads_one_row_and_a_qualified_alias_is_unsupported() } other => panic!("{other:?}"), } - let qualified = relation(RelType::SubqueryAlias(Box::new(pb::SubqueryAlias { +} + +fn aliased(qualifier: &[&str], alias: &str) -> Relation { + relation(RelType::SubqueryAlias(Box::new(pb::SubqueryAlias { input: input(), - alias: "x".into(), - qualifier: vec!["db".into()], - }))); - assert!(message(&qualified).contains("SubqueryAlias.qualifier")); + alias: alias.into(), + qualifier: qualifier.iter().map(|q| (*q).to_owned()).collect(), + }))) +} + +#[test] +fn a_subquery_alias_qualifier_prefixes_the_alias_parts() { + let sep = krishiv_plan::ALIAS_PART_SEP; + let PlanNode::SubqueryAlias { alias, .. } = node(&aliased(&["db", "s"], "t")) else { + panic!("not an alias") + }; + assert_eq!(alias, format!("db{sep}s{sep}t")); + let PlanNode::SubqueryAlias { alias, .. } = node(&aliased(&[], "t")) else { + panic!("not an alias") + }; + assert_eq!(alias, "t"); + // Printed dotted, as the IR prints a multi-part alias. + assert_eq!( + show(&aliased(&["db"], "t")), + "SubqueryAlias db.t\n+- Read t\n" + ); +} + +#[test] +fn an_alias_part_holding_the_separator_is_refused() { + let sep = krishiv_plan::ALIAS_PART_SEP; + for r in [ + aliased(&[&format!("a{sep}b")], "t"), + aliased(&["db"], &format!("t{sep}u")), + aliased(&[], &format!("t{sep}u")), + ] { + let e = fails(translate(&r)); + assert_eq!(e.error_class().name(), "KRISHIV_UNSUPPORTED_FEATURE"); + assert!( + e.message().contains("SubqueryAlias.qualifier"), + "{}", + e.message() + ); + } +} + +fn read_with_schema(schema: &str) -> Relation { + relation(RelType::Read(pb::Read { + is_streaming: false, + read_type: Some(ReadType::DataSource(read::DataSource { + format: Some("csv".into()), + paths: vec!["/a.csv".into()], + schema: Some(schema.into()), + ..read::DataSource::default() + })), + })) +} + +#[test] +fn a_data_source_schema_is_the_reads_schema() { + let ddl = node(&read_with_schema("a INT NOT NULL, b STRING")); + let PlanNode::Read(ReadSource::Files { schema, .. }) = &ddl else { + panic!("{ddl:?}") + }; + let schema = schema.as_ref().expect("a schema"); + assert_eq!(schema.fields().len(), 2); + assert!(!schema.fields()[0].nullable()); + // JSON is the fallback. + let json = node(&read_with_schema( + r#"{"type":"struct","fields":[{"name":"a","type":"integer","nullable":false,"metadata":{}},{"name":"b","type":"string","nullable":true,"metadata":{}}]}"#, + )); + let PlanNode::Read(ReadSource::Files { + schema: from_json, .. + }) = &json + else { + panic!("{json:?}") + }; + assert_eq!(from_json.as_ref(), Some(schema)); + // A schema that is not a struct (JSON of a bare type) is an invalid plan. + assert_eq!( + fails(translate(&read_with_schema("\"integer\""))) + .error_class() + .name(), + "CONNECT_INVALID_PLAN.INVALID_SCHEMA_NON_STRUCT_TYPE" + ); + // Neither DDL nor JSON: a parse error, and the text is not echoed. + let e = fails(translate(&read_with_schema("CANARY nonsense nonsense"))); + assert!(!e.message().contains("CANARY"), "{}", e.message()); } #[test] @@ -797,10 +871,6 @@ fn unsupported_relations_are_named() { (RelType::ShowString(Box::default()), "SHOW_STRING(20)"), (RelType::ToDf(Box::default()), "TO_DF(18)"), (RelType::Deduplicate(Box::default()), "DEDUPLICATE(14)"), - ( - RelType::LocalRelation(Default::default()), - "LOCAL_RELATION(11)", - ), (RelType::Drop(Box::default()), "DROP(21)"), (RelType::WithRelations(Box::default()), "WITH_RELATIONS(41)"), ] { @@ -874,15 +944,19 @@ fn secrets_in_read_options_reach_no_printer_key_or_error() { } let e = fails(translate(&streaming)); assert!(!e.to_string().contains(CANARY)); - let mut with_schema = r; + let mut with_predicates = r; if let Some(RelType::Read(pb::Read { read_type: Some(ReadType::DataSource(d)), .. - })) = &mut with_schema.rel_type + })) = &mut with_predicates.rel_type { - d.schema = Some("a INT".into()); + d.predicates = vec!["a < 1".into()]; } - assert!(!fails(translate(&with_schema)).to_string().contains(CANARY)); + assert!( + !fails(translate(&with_predicates)) + .to_string() + .contains(CANARY) + ); // The printed plan still names the options. assert!( plan.to_string().contains("password=*********(redacted)"), diff --git a/crates/krishiv-connect/src/translate/schema.rs b/crates/krishiv-connect/src/translate/schema.rs new file mode 100644 index 000000000..8c410c63e --- /dev/null +++ b/crates/krishiv-connect/src/translate/schema.rs @@ -0,0 +1,77 @@ +// SPDX-License-Identifier: Apache-2.0 +//! Schema texts a Connect plan carries (`LocalRelation.schema`, `Read.DataSource.schema`): DDL or +//! JSON, parsed as Spark's `SparkConnectPlanner` does (v4.3.0). + +use krishiv_common::error::{ConnectInvalidPlanInvalidSchemaNonStructType, KrishivError}; +use krishiv_types::{DataType, StructField, StructType}; + +/// The planner's `parseDatatypeString` (`SparkConnectPlanner.scala:1511-1528`): +/// a table schema (`a INT, b STRING`), else a data type (`int`), else the +/// text as the fields of a struct; the first failure is the error. +/// +/// Krishiv parses the table schema as the fields of a `STRUCT<...>`, so +/// clauses only a table schema takes (`DEFAULT`, `GENERATED`) are refused, +/// and the error is the data type's (Spark reports the table schema's): a +/// recorded difference. +/// +/// # Errors +/// The parse error of the text as a data type. +pub fn parse_datatype_string(text: &str) -> Result { + if let Ok(t) = krishiv_types::parse_data_type(&format!("STRUCT<{text}>")) { + return Ok(t); + } + match krishiv_types::parse_data_type(text) { + Ok(t) => Ok(t), + Err(e) => { + krishiv_types::parse_data_type(&format!("STRUCT<{}>", text.trim())).map_err(|_| e) + } + } +} + +/// `DataType.parseTypeWithFallback`: `parser`'s result, else `fallback`'s; when both fail, the +/// first error. +fn with_fallback( + text: &str, + parser: impl FnOnce(&str) -> Result, + fallback: impl FnOnce(&str) -> Result, +) -> Result { + parser(text).or_else(|first| fallback(text).map_err(|_| first)) +} + +/// A `LocalRelation` schema (`SparkConnectPlanner.transformLocalRelation`): DDL by +/// [`parse_datatype_string`], else JSON, and a type that is not a struct wrapped as +/// `struct` (`toStructTypeOrWrap`). +/// +/// # Errors +/// The parse error of the text as DDL when it is not JSON either. +pub fn parse_local_schema(text: &str) -> Result { + Ok( + match with_fallback(text, parse_datatype_string, DataType::from_json)? { + DataType::Struct(s) => s, + other => StructType::new(vec![StructField::new("value", other, true)]), + }, + ) +} + +/// A data source's user schema (`parseSchema`, `SparkConnectPlanner.scala:1682`): a table schema +/// by `StructType.fromDDL`, else JSON. +/// +/// # Errors +/// The parse error of the DDL when the text is not JSON either; +/// `CONNECT_INVALID_PLAN.INVALID_SCHEMA_NON_STRUCT_TYPE` for JSON of a type that is not a struct +/// (golden-pending: the class is the registry's Connect one for +/// `invalidSchemaStringNonStructType`, its use unverified against a real run). +pub fn parse_table_schema(text: &str) -> Result { + let ddl = |t: &str| krishiv_types::parse_data_type(&format!("STRUCT<{t}>")); + match with_fallback(text, ddl, DataType::from_json)? { + DataType::Struct(s) => Ok(s), + other => Err(ConnectInvalidPlanInvalidSchemaNonStructType { + input_schema: format!("\"{text}\""), + data_type: format!("\"{}\"", other.sql()), + } + .into()), + } +} + +#[cfg(test)] +mod tests; diff --git a/crates/krishiv-connect/src/translate/schema/tests.rs b/crates/krishiv-connect/src/translate/schema/tests.rs new file mode 100644 index 000000000..5d2eef1b5 --- /dev/null +++ b/crates/krishiv-connect/src/translate/schema/tests.rs @@ -0,0 +1,55 @@ +// SPDX-License-Identifier: Apache-2.0 +//! Tests of the schema texts. + +use krishiv_types::DataType; + +use super::{parse_datatype_string, parse_local_schema, parse_table_schema}; + +#[test] +fn a_local_schema_is_ddl_a_type_or_json_and_a_type_is_wrapped() { + assert_eq!( + parse_local_schema("a INT, b STRING NOT NULL").unwrap(), + parse_local_schema( + r#"{"type":"struct","fields":[{"name":"a","type":"integer","nullable":true,"metadata":{}},{"name":"b","type":"string","nullable":false,"metadata":{}}]}"# + ) + .unwrap() + ); + // A bare type becomes `struct`. + let wrapped = parse_local_schema("int").unwrap(); + assert_eq!(wrapped.fields().len(), 1); + assert_eq!(wrapped.fields()[0].name(), "value"); + let json_type = parse_local_schema("\"long\"").unwrap(); + assert_eq!(json_type.fields()[0].name(), "value"); + // Neither DDL nor JSON: the DDL error. + assert!(parse_local_schema("a nonsense_type").is_err()); +} + +#[test] +fn a_data_source_schema_must_be_a_struct() { + let s = parse_table_schema("a INT, b ARRAY").unwrap(); + assert_eq!( + DataType::Struct(s).simple_string(), + "struct>" + ); + let json = parse_table_schema( + r#"{"type":"struct","fields":[{"name":"a","type":"integer","nullable":false,"metadata":{}}]}"#, + ) + .unwrap(); + assert!(!json.fields()[0].nullable()); + let e = parse_table_schema("\"integer\"").unwrap_err(); + assert_eq!( + e.error_class().name(), + "CONNECT_INVALID_PLAN.INVALID_SCHEMA_NON_STRUCT_TYPE" + ); + // Not DDL, not JSON: the DDL error, never an echo of the text. + assert!(parse_table_schema("a b c d").is_err()); +} + +#[test] +fn datatype_strings_try_schema_type_then_struct() { + assert_eq!( + parse_datatype_string(" a: int ").unwrap().simple_string(), + "struct" + ); + assert_eq!(parse_datatype_string("int").unwrap().simple_string(), "int"); +} diff --git a/crates/krishiv-connect/src/translate/tests.rs b/crates/krishiv-connect/src/translate/tests.rs index ed15bc853..91662a0f1 100644 --- a/crates/krishiv-connect/src/translate/tests.rs +++ b/crates/krishiv-connect/src/translate/tests.rs @@ -67,7 +67,7 @@ fn commands_compressed_plans_and_empty_plans_are_not_relations() { ); let command = pb::Plan { op_type: Some(OpType::Command(Box::new(pb::Command { - command_type: Some(pb::command::CommandType::SqlCommand(Box::default())), + command_type: Some(pb::command::CommandType::GetResourcesCommand(Box::default())), }))), }; let e = fails(translate_plan(&command, &conf())); @@ -75,7 +75,11 @@ fn commands_compressed_plans_and_empty_plans_are_not_relations() { e.error_class(), ErrorClass::ConnectInvalidPlanInvalidOneOfFieldNotSupported ); - assert!(e.message().contains("SQL_COMMAND(5)"), "{}", e.message()); + assert!( + e.message().contains("GET_RESOURCES_COMMAND(8)"), + "{}", + e.message() + ); let compressed = pb::Plan { op_type: Some(OpType::CompressedOperation(Box::default())), }; @@ -294,3 +298,73 @@ fn debug_of_a_write_hides_a_path_held_by_a_oneof() { assert!(!text.contains("Canary-4d1"), "{text}"); assert!(text.contains("Path()"), "{text}"); } + +#[test] +fn a_sql_command_translates_as_the_sql_relation_of_a_query() { + #[expect(deprecated, reason = "the field every client version sends")] + let command = pb::SqlCommand { + sql: "SELECT 1".to_owned(), + pos_args: vec![pb::expression::Literal { + literal_type: Some(pb::expression::literal::LiteralType::Integer(5)), + data_type: None, + }], + ..Default::default() + }; + let relation = sql_command_relation(&command).expect("a relation"); + let Some(pb::relation::RelType::Sql(sql)) = &relation.rel_type else { + panic!("a SQL relation: {relation:?}"); + }; + assert_eq!(sql.query, "SELECT 1"); + assert!(sql_has_parameters(sql)); + assert!(!sql_has_parameters(&pb::Sql::default())); + let plan = pb::Plan { + op_type: Some(OpType::Command(Box::new(pb::Command { + command_type: Some(pb::command::CommandType::SqlCommand(Box::new(command))), + }))), + }; + assert_eq!( + translate_plan(&plan, &conf()) + .expect("translates") + .to_string(), + translate_relation(&relation, &conf()) + .expect("translates") + .to_string() + ); + // `input` (PySpark 4.x) wins and is answered as it came, plan id included. + let input = pb::Relation { + common: Some(pb::RelationCommon { + plan_id: Some(9), + ..Default::default() + }), + rel_type: Some(pb::relation::RelType::Sql(pb::Sql { + query: "SELECT 2".to_owned(), + ..Default::default() + })), + }; + let with_input = pb::SqlCommand { + input: Some(input.clone()), + ..Default::default() + }; + assert_eq!(sql_command_relation(&with_input).expect("input"), input); + // The `WithRelations` form is refused as an unsupported field. + let with_relations = pb::SqlCommand { + input: Some(pb::Relation { + common: None, + rel_type: Some(pb::relation::RelType::WithRelations(Box::default())), + }), + ..Default::default() + }; + let e = sql_command_relation(&with_relations).expect_err("refused"); + assert_eq!(e.error_class(), ErrorClass::KrishivUnsupportedFeature); + assert!(e.message().contains("SqlCommand.input"), "{}", e.message()); + // Any other relation is Spark's `SQL_COMMAND_EXPECTS_SQL_OR_WITH_RELATIONS`. + let other = pb::SqlCommand { + input: Some(pb::Relation::default()), + ..Default::default() + }; + let e = sql_command_relation(&other).expect_err("refused"); + assert_eq!( + e.error_class(), + ErrorClass::ConnectInvalidPlanSqlCommandExpectsSqlOrWithRelations + ); +} diff --git a/crates/krishiv-connect/tests/goldens.rs b/crates/krishiv-connect/tests/goldens.rs index 5ff13c539..90a195025 100644 --- a/crates/krishiv-connect/tests/goldens.rs +++ b/crates/krishiv-connect/tests/goldens.rs @@ -28,19 +28,6 @@ const UNSUPPORTED: &[(&str, ErrorClass)] = &[ "distinct", ErrorClass::ConnectInvalidPlanInvalidOneOfFieldNotSupported, ), - ( - "filter_expression_string", - ErrorClass::ConnectInvalidPlanInvalidOneOfFieldNotSupported, - ), - ( - "literals_array", - ErrorClass::ConnectInvalidPlanInvalidOneOfFieldNotSupported, - ), - ("read_schema_string", ErrorClass::KrishivUnsupportedFeature), - ( - "select_expr", - ErrorClass::ConnectInvalidPlanInvalidOneOfFieldNotSupported, - ), ( "show_string", ErrorClass::ConnectInvalidPlanInvalidOneOfFieldNotSupported, diff --git a/crates/krishiv-connect/tests/snapshots/filter_expression_string.snap b/crates/krishiv-connect/tests/snapshots/filter_expression_string.snap index aed823ca2..e3e5a452b 100644 --- a/crates/krishiv-connect/tests/snapshots/filter_expression_string.snap +++ b/crates/krishiv-connect/tests/snapshots/filter_expression_string.snap @@ -1,4 +1,8 @@ --- source: crates/krishiv-connect/tests/goldens.rs --- -[CONNECT_INVALID_PLAN.INVALID_ONE_OF_FIELD_NOT_SUPPORTED] The Spark Connect plan is invalid. This oneOf field message in spark.connect.Expression is not supported: EXPRESSION_STRING(4) SQLSTATE: 56K00 +Filter ((a > 1) and (b < 2)) ++- Read parquet [/data/t.parquet] +-- plan ids -- +Filter plan_id=1 + Read plan_id=0 diff --git a/crates/krishiv-connect/tests/snapshots/literals_array.snap b/crates/krishiv-connect/tests/snapshots/literals_array.snap index 5cc5ef446..9358ec2d2 100644 --- a/crates/krishiv-connect/tests/snapshots/literals_array.snap +++ b/crates/krishiv-connect/tests/snapshots/literals_array.snap @@ -1,4 +1,8 @@ --- source: crates/krishiv-connect/tests/goldens.rs --- -[CONNECT_INVALID_PLAN.INVALID_ONE_OF_FIELD_NOT_SUPPORTED] The Spark Connect plan is invalid. This oneOf field message in spark.connect.Expression.Literal is not supported: ARRAY(22) SQLSTATE: 56K00 +Project [array(1, 2), ARRAY(1, 2, 3)] ++- Read parquet [/data/t.parquet] +-- plan ids -- +Project plan_id=1 + Read plan_id=0 diff --git a/crates/krishiv-connect/tests/snapshots/read_schema_string.snap b/crates/krishiv-connect/tests/snapshots/read_schema_string.snap index d5bbc181e..a5a69987c 100644 --- a/crates/krishiv-connect/tests/snapshots/read_schema_string.snap +++ b/crates/krishiv-connect/tests/snapshots/read_schema_string.snap @@ -1,4 +1,6 @@ --- source: crates/krishiv-connect/tests/goldens.rs --- -[KRISHIV_UNSUPPORTED_FEATURE] The feature reads with a user-specified schema (spark.connect.Read.DataSource.schema) is not supported by Krishiv yet. SQLSTATE: 0A000 +Read parquet [/data/p] schema=a INT, b STRING +-- plan ids -- +Read plan_id=0 diff --git a/crates/krishiv-connect/tests/snapshots/select_expr.snap b/crates/krishiv-connect/tests/snapshots/select_expr.snap index aed823ca2..f3c4b01d1 100644 --- a/crates/krishiv-connect/tests/snapshots/select_expr.snap +++ b/crates/krishiv-connect/tests/snapshots/select_expr.snap @@ -1,4 +1,8 @@ --- source: crates/krishiv-connect/tests/goldens.rs --- -[CONNECT_INVALID_PLAN.INVALID_ONE_OF_FIELD_NOT_SUPPORTED] The Spark Connect plan is invalid. This oneOf field message in spark.connect.Expression is not supported: EXPRESSION_STRING(4) SQLSTATE: 56K00 +Project [(a + 1) AS x] ++- Read parquet [/data/t.parquet] +-- plan ids -- +Project plan_id=1 + Read plan_id=0 diff --git a/crates/krishiv-session/src/lib.rs b/crates/krishiv-session/src/lib.rs index 475919e82..08c686f43 100644 --- a/crates/krishiv-session/src/lib.rs +++ b/crates/krishiv-session/src/lib.rs @@ -137,6 +137,8 @@ pub use krishiv_udf_sdk::{ pub use profile::{ExplainVerbose, OptimizerStep, VERBOSE_FORMAT, VERBOSE_VERSION}; pub use python_udfs::{ANONYMOUS_VERSION_BIT, PythonUdfDef, PythonUdfEntry, PythonUdfLibrary}; pub use runtime::{EngineRuntime, EngineRuntimeBuilder}; +#[cfg(feature = "unstable-ir")] +pub use session::SqlOutcome; pub use session::{RecordBatchStream, Session, SessionBuilder}; /// Native Rust scalar UDFs: the [`ScalarUdf`] trait and its boundary types, /// for [`Session::register_udf`]. diff --git a/crates/krishiv-session/src/session.rs b/crates/krishiv-session/src/session.rs index 33eccf9fa..dff911bed 100644 --- a/crates/krishiv-session/src/session.rs +++ b/crates/krishiv-session/src/session.rs @@ -64,6 +64,8 @@ use crate::variables::SessionVariables; mod context; mod lifecycle; +#[cfg(feature = "unstable-ir")] +pub use lifecycle::SqlOutcome; #[cfg(test)] mod run_task; @@ -491,6 +493,11 @@ impl Session { /// Spark's parse and analysis errors; `KRISHIV_UNSUPPORTED_FEATURE` /// for statements other than queries and commands. pub async fn sql(&self, query: &str) -> Result { + self.sql_kind(query).await.map(|(df, _)| df) + } + + /// [`Self::sql`] and whether the statement was a command (not a query). + pub(crate) async fn sql_kind(&self, query: &str) -> Result<(DataFrame, bool), KrishivError> { self.check_open()?; let (session, text) = (self.clone(), query.to_owned()); let row: Arc>> = Arc::default(); @@ -505,7 +512,7 @@ impl Session { .sql_here_marked(&text, marker) .instrument(span) .await - .map(|df| df.with_trace(trace)) + .map(|(df, command)| (df.with_trace(trace), command)) }) .await; // A statement refused before it ran (parse, analysis, access) is a @@ -552,7 +559,9 @@ impl Session { spans: Vec, ) -> Result { // Boxed: one more nested future level overflows the compiler's query depth in the binaries. - Box::pin(self.sql_here_inner(query, Arc::default(), spans)).await + Box::pin(self.sql_here_inner(query, Arc::default(), spans)) + .await + .map(|(df, _)| df) } /// Runs `query` on the calling task as `sql` would, with `row` as the statement's slot for @@ -561,7 +570,7 @@ impl Session { &self, query: &str, row: Arc>>, - ) -> Result { + ) -> Result<(DataFrame, bool), KrishivError> { // Boxed: one more nested future level overflows the compiler's query depth in the binaries. Box::pin(self.sql_here_inner(query, row, Vec::new())).await } @@ -571,7 +580,7 @@ impl Session { query: &str, row: Arc>>, spans: Vec, - ) -> Result { + ) -> Result<(DataFrame, bool), KrishivError> { let mut env = self.query_env(); env.row = row; let written = query; @@ -601,11 +610,12 @@ impl Session { let class = krishiv_sql::classify(&statement); self.authorize_statement(&statement, class, &env).await?; env.statement_class = Some(class.name()); + let command = !matches!(statement, Statement::Query(_)); let df = self.run_statement(statement, env, query, &redacted).await?; // Whatever frame the statement returns (a command's result too) shows the statement's // text and class when it runs. df.with_statement(&redacted, class.name()); - Ok(df) + Ok((df, command)) } /// Runs the parsed, authorized `statement`. @@ -868,7 +878,15 @@ impl Session { Ok(Column::new((*lowered).clone())) } + /// This session's parser options, from its configuration. Public under `unstable-ir`, for a + /// Connect server that parses `Expression.expression_string` while it translates a plan. + #[cfg(feature = "unstable-ir")] + pub fn parser_options(&self) -> ParserOptions { + parser_options(&self.state.conf.snapshot()) + } + /// This session's parser options, from its configuration. + #[cfg(not(feature = "unstable-ir"))] pub(crate) fn parser_options(&self) -> ParserOptions { parser_options(&self.state.conf.snapshot()) } diff --git a/crates/krishiv-session/src/session/lifecycle.rs b/crates/krishiv-session/src/session/lifecycle.rs index 00b27ca24..81ff80e99 100644 --- a/crates/krishiv-session/src/session/lifecycle.rs +++ b/crates/krishiv-session/src/session/lifecycle.rs @@ -29,6 +29,21 @@ use crate::sql_functions::SqlFunctions; use crate::udfs::SessionFunctions; use crate::variables::SessionVariables; +/// What [`Session::sql_command`] ran. +#[cfg(feature = "unstable-ir")] +#[expect( + missing_debug_implementations, + reason = "a `DataFrame` has no `Debug`; neither has what holds one" +)] +#[non_exhaustive] +pub enum SqlOutcome { + /// A query: analyzed, not run; its rows come from the frame. + Query(DataFrame), + /// A command (`SET`, `CREATE VIEW`, `INSERT`, `SHOW`, ...), already run; the frame holds its + /// result rows. + Command(DataFrame), +} + impl Session { /// A session on `runtime` with configuration `config`. pub(super) fn on_runtime( @@ -422,6 +437,45 @@ impl Session { self.with_options(options)?.sql(query).await } + /// Runs `query` as [`Self::sql`] does and says what it was: a lazy query, or a command (any + /// statement other than a query, which `sql` runs at once, as Spark's eager `Command`s). + /// + /// Spark Connect's `SqlCommand` needs the difference: a command's result is returned to the + /// client as a local relation, a query's as the `SQL` relation to send again. + /// + /// ``` + /// # let session = krishiv_session::Session::new()?; + /// use krishiv_session::SqlOutcome; + /// let q = futures::executor::block_on(session.sql_command("SELECT 1"))?; + /// assert!(matches!(q, SqlOutcome::Query(_))); + /// let c = futures::executor::block_on(session.sql_command("SET a.b = 1"))?; + /// assert!(matches!(c, SqlOutcome::Command(_))); + /// # Ok::<(), krishiv_common::error::KrishivError>(()) + /// ``` + /// + /// # Errors + /// As [`Self::sql`]. + #[cfg(feature = "unstable-ir")] + pub async fn sql_command(&self, query: &str) -> Result { + Ok(match self.sql_kind(query).await? { + (df, false) => crate::SqlOutcome::Query(df), + (df, true) => crate::SqlOutcome::Command(df), + }) + } + + /// As [`Self::sql_command`], with per-statement options (as [`Self::sql_with`]). + /// + /// # Errors + /// As [`Self::sql_with`]. + #[cfg(feature = "unstable-ir")] + pub async fn sql_command_with( + &self, + query: &str, + options: crate::QueryOptions, + ) -> Result { + self.with_options(options)?.sql_command(query).await + } + /// A handle of this session that runs its queries under `options` (the /// operation id taken now, held while the handle or a clone lives). pub(crate) fn with_options( diff --git a/docs/reference/connect/coverage.md b/docs/reference/connect/coverage.md index f8e702aba..3aa7582a0 100644 --- a/docs/reference/connect/coverage.md +++ b/docs/reference/connect/coverage.md @@ -5,11 +5,11 @@ Every arm of the `oneof` fields the Connect translator dispatches on, for the ve ## Relation.rel_type -62 arms: 7 done, 6 partial, 45 missing, 4 out of scope. +62 arms: 9 done, 6 partial, 43 missing, 4 out of scope. | # | Arm | Status | Refused fields | IR target | Task | Group | |---:|---|---|---|---|---|---:| -| 2 | read | partial | DataSource.schema, predicates (JDBC) and is_streaming | Read(ReadSource::{Table,Files}); user schema on ReadSource::Files; predicates → JDBC partitions (refuse as Spark does for non-JDBC formats: PREDICATES_NOT_SUPPORTED_FOR_DATA_SOURCE) | P3-T012 | 1 | +| 2 | read | partial | DataSource.predicates (JDBC) and is_streaming | Read(ReadSource::{Table,Files}); DataSource.schema → ReadSource::Files schema (DDL or JSON, a non-struct is INVALID_SCHEMA_NON_STRUCT_TYPE); predicates → JDBC partitions (refuse as Spark does for non-JDBC formats: PREDICATES_NOT_SUPPORTED_FOR_DATA_SOURCE) | P3-T012 | 1 | | 3 | project | done | | Project; missing input → OneRowRelation | P3-T012 | 1 | | 4 | filter | done | | Filter | P2-T091a | 0 | | 5 | join | partial | join_data_type (joinWith, typed) | Join; join_data_type is Scala joinWith only → unsupported row | P3-T012 | 1 | @@ -18,12 +18,12 @@ Every arm of the `oneof` fields the Connect translator dispatches on, for the ve | 8 | limit | done | | Limit | P2-T091a | 0 | | 9 | aggregate | partial | GROUP_TYPE_PIVOT | Aggregate; pivot → Pivot (exists) | P3-T012 | 2 | | 10 | sql | partial | DataFrame arguments (with_relations) | Sql; the with_relations form (#41) for DataFrame arguments | P3-T012 | 4 | -| 11 | local_relation | missing | | InlineTable from the Arrow IPC data plus the DDL or JSON schema (createDataFrame); a size limit on the data | P3-T012 | 1 | +| 11 | local_relation | done | | Read(ReadSource::Local) from the Arrow IPC data plus the DDL or JSON schema (createDataFrame), built by the host before translation (TranslateConf::local_relations); no data and no schema is SCHEMA_REQUIRED_FOR_LOCAL_RELATION; the inline size is bounded by the gRPC message limit | P3-T012a | 1 | | 12 | sample | missing | | new Sample (Bernoulli/Poisson with seed, Spark's XORShiftRandom order) | P3-T012 | 3 | | 13 | offset | done | | Offset | P2-T091a | 0 | | 14 | deduplicate | missing | | lowering: Distinct (all columns) or Aggregate with first per non-key column; within_watermark → P4 | P3-T012 | 3 | | 15 | range | done | | Range | P2-T091a | 0 | -| 16 | subquery_alias | partial | qualifier | SubqueryAlias with the qualifier | P3-T012 | 1 | +| 16 | subquery_alias | done | | SubqueryAlias; the qualifier and the alias are the alias parts joined by ALIAS_PART_SEP | P3-T012 | 1 | | 17 | repartition | missing | | Repartition (round robin / coalesce) | P3-T012 | 3 | | 18 | to_df | missing | | lowering: Project of Alias (positional) | P3-T012 | 4 | | 19 | with_columns_renamed | missing | | lowering: Project | P3-T012 | 4 | @@ -43,7 +43,7 @@ Every arm of the `oneof` fields the Connect translator dispatches on, for the ve | 33 | with_watermark | out of scope | | P4 | | 0 | | 34 | apply_in_pandas_with_state | out of scope | | P4 | | 0 | | 35 | html_string | missing | | as show_string, HTML form | P3-T012 | 4 | -| 36 | cached_local_relation | missing | | InlineTable from the artifact cache (AddArtifacts cache/, T011); errors NOT_FOUND_CACHED_LOCAL_RELATION | P3-T083 | 7 | +| 36 | cached_local_relation | partial | hashes in the artifact store (AddArtifacts cache/ is not stored yet: every hash is NOT_FOUND_CACHED_LOCAL_RELATION) | Read(ReadSource::Local) through the same host request as local_relation; the artifact lookup is the host's | P3-T083 | 7 | | 37 | cached_remote_relation | missing | | the session's cached-DataFrame registry (T083) | P3-T083 | 7 | | 38 | common_inline_user_defined_table_function | missing | | TableFunction with a Python UDTF (T014) | P3-T014 | 6 | | 39 | as_of_join | missing | | new AsOfJoin, shared with the ASOF JOIN SQL syntax | P3-T012 | 8 | @@ -74,14 +74,14 @@ Every arm of the `oneof` fields the Connect translator dispatches on, for the ve ## Expression.expr_type -22 arms: 3 done, 5 partial, 13 missing, 1 out of scope. +22 arms: 4 done, 5 partial, 12 missing, 1 out of scope. | # | Arm | Status | Refused fields | IR target | Task | Group | |---:|---|---|---|---|---|---:| -| 1 | literal | partial | Array, Map, Struct, SpecializedArray, Time | Literal/TypedLiteral; nested literals after the IR literal supports them; nanosecond timestamps are refused until the types exist | P3-T012 | 1 | +| 1 | literal | partial | Time | Literal/TypedLiteral; nested literals are Literal::{Array,KeyValueMap,Struct}; nanosecond timestamps are refused until the types exist | P3-T012 | 1 | | 2 | unresolved_attribute | partial | is_metadata_column | Column; metadata columns (_metadata) → the scan's metadata columns | P3-T012 | 1 | | 3 | unresolved_function | partial | is_user_defined_function (qualified names) and is_internal | Function with qualifier; internal functions map to Spark's internal names where Krishiv has them, otherwise 56K00-style refusal naming the function | P3-T012 | 1 | -| 4 | expression_string | missing | | the SQL expression parser + analyzer lowering that select_expr/expr() use (krishiv-session) | P3-T012 | 1 | +| 4 | expression_string | done | | krishiv_sql::parse_and_lower_named_expression, as Session::select_expr and Session::expr (parser options from TranslateConf) | P3-T012a | 1 | | 5 | unresolved_star | done | | Star | P2-T091a | 0 | | 6 | alias | partial | metadata and several names | Alias; several names → MultiAlias (exists) | P3-T012 | 1 | | 7 | cast | done | | Cast | P2-T091a | 0 | @@ -103,7 +103,7 @@ Every arm of the `oneof` fields the Connect translator dispatches on, for the ve ## Command.command_type -20 arms: 0 done, 0 partial, 14 missing, 6 out of scope. +20 arms: 0 done, 1 partial, 13 missing, 6 out of scope. | # | Arm | Status | Refused fields | IR target | Task | Group | |---:|---|---|---|---|---|---:| @@ -111,7 +111,7 @@ Every arm of the `oneof` fields the Connect translator dispatches on, for the ve | 2 | write_operation | missing | | DataFrameWriter V1: krishiv-session save/save_as_table/insert_into (dataframe.rs), semantics T023 | P3-T023 | 11 | | 3 | create_dataframe_view | missing | | create_or_replace_temp_view (exists), global temp views | P3-T012 | 4 | | 4 | write_operation_v2 | missing | | DataFrameWriterV2 (create/replace/append/overwrite/overwrite_partitions exist in dataframe.rs), T023 | P3-T023 | 11 | -| 5 | sql_command | missing | | spark.sql() sends this. Run commands at once; for queries return sql_command_result.relation (Spark returns a SQL or LocalRelation relation the client re-sends). Without it no spark.sql test passes | P3-T012 | 1 | +| 5 | sql_command | partial | input of WithRelations form (group 4; the plain SQL input is supported); named and positional arguments (the analyzer refuses SQL parameters) | spark.sql() sends this. Run commands at once; for queries return sql_command_result.relation (Spark returns a SQL or LocalRelation relation the client re-sends). Translated here as the SQL relation of a query; the server runs commands (execute.rs) | P3-T012 | 1 | | 6 | write_stream_operation_start | out of scope | | P4 | | 0 | | 7 | streaming_query_command | out of scope | | P4 | | 0 | | 8 | get_resources_command | missing | | an empty resource map (single node, no executors) | P3-T012 | 4 | @@ -173,7 +173,7 @@ Every arm of the `oneof` fields the Connect translator dispatches on, for the ve ## Expression.Literal.literal_type -24 arms: 17 done, 0 partial, 7 missing, 0 out of scope. +24 arms: 21 done, 0 partial, 3 missing, 0 out of scope. | # | Arm | Status | Refused fields | IR target | Task | Group | |---:|---|---|---|---|---|---:| @@ -194,11 +194,11 @@ Every arm of the `oneof` fields the Connect translator dispatches on, for the ve | 19 | calendar_interval | done | | Literal / TypedLiteral | P2-T091a | 1 | | 20 | year_month_interval | done | | Literal / TypedLiteral | P2-T091a | 1 | | 21 | day_time_interval | done | | Literal / TypedLiteral | P2-T091a | 1 | -| 22 | array | missing | | nested and TIME literals wait for IR values | P3-T012 | 1 | -| 23 | map | missing | | nested and TIME literals wait for IR values | P3-T012 | 1 | -| 24 | struct | missing | | nested and TIME literals wait for IR values | P3-T012 | 1 | -| 25 | specialized_array | missing | | nested and TIME literals wait for IR values | P3-T012 | 1 | -| 26 | time | missing | | nested and TIME literals wait for IR values | P3-T012 | 1 | +| 22 | array | done | | Literal::{Array,KeyValueMap,Struct}; the type is the root data_type, else the message's own type fields | P3-T012a | 1 | +| 23 | map | done | | Literal::{Array,KeyValueMap,Struct}; the type is the root data_type, else the message's own type fields | P3-T012a | 1 | +| 24 | struct | done | | Literal::{Array,KeyValueMap,Struct}; the type is the root data_type, else the message's own type fields | P3-T012a | 1 | +| 25 | specialized_array | done | | Literal::{Array,KeyValueMap,Struct}; the type is the root data_type, else the message's own type fields | P3-T012a | 1 | +| 26 | time | missing | | TIME literals wait for an IR value | P3-T012 | 1 | | 29 | timestamp_ntz_nanos | missing | | no IR type for nanosecond timestamps | P3-T012 | 1 | | 30 | timestamp_ltz_nanos | missing | | no IR type for nanosecond timestamps | P3-T012 | 1 | diff --git a/tests/connect/coverage.json b/tests/connect/coverage.json index a48c98101..aa74c3966 100644 --- a/tests/connect/coverage.json +++ b/tests/connect/coverage.json @@ -1,5 +1,5 @@ [ - {"oneof": "Relation.rel_type", "field": "read", "number": 2, "status": "partial", "refused": "DataSource.schema, predicates (JDBC) and is_streaming", "target": "Read(ReadSource::{Table,Files}); user schema on ReadSource::Files; predicates → JDBC partitions (refuse as Spark does for non-JDBC formats: PREDICATES_NOT_SUPPORTED_FOR_DATA_SOURCE)", "task": "P3-T012", "group": 1}, + {"oneof": "Relation.rel_type", "field": "read", "number": 2, "status": "partial", "refused": "DataSource.predicates (JDBC) and is_streaming", "target": "Read(ReadSource::{Table,Files}); DataSource.schema → ReadSource::Files schema (DDL or JSON, a non-struct is INVALID_SCHEMA_NON_STRUCT_TYPE); predicates → JDBC partitions (refuse as Spark does for non-JDBC formats: PREDICATES_NOT_SUPPORTED_FOR_DATA_SOURCE)", "task": "P3-T012", "group": 1}, {"oneof": "Relation.rel_type", "field": "project", "number": 3, "status": "done", "refused": "", "target": "Project; missing input → OneRowRelation", "task": "P3-T012", "group": 1}, {"oneof": "Relation.rel_type", "field": "filter", "number": 4, "status": "done", "refused": "", "target": "Filter", "task": "P2-T091a", "group": 0}, {"oneof": "Relation.rel_type", "field": "join", "number": 5, "status": "partial", "refused": "join_data_type (joinWith, typed)", "target": "Join; join_data_type is Scala joinWith only → unsupported row", "task": "P3-T012", "group": 1}, @@ -8,12 +8,12 @@ {"oneof": "Relation.rel_type", "field": "limit", "number": 8, "status": "done", "refused": "", "target": "Limit", "task": "P2-T091a", "group": 0}, {"oneof": "Relation.rel_type", "field": "aggregate", "number": 9, "status": "partial", "refused": "GROUP_TYPE_PIVOT", "target": "Aggregate; pivot → Pivot (exists)", "task": "P3-T012", "group": 2}, {"oneof": "Relation.rel_type", "field": "sql", "number": 10, "status": "partial", "refused": "DataFrame arguments (with_relations)", "target": "Sql; the with_relations form (#41) for DataFrame arguments", "task": "P3-T012", "group": 4}, - {"oneof": "Relation.rel_type", "field": "local_relation", "number": 11, "status": "missing", "refused": "", "target": "InlineTable from the Arrow IPC data plus the DDL or JSON schema (createDataFrame); a size limit on the data", "task": "P3-T012", "group": 1}, + {"oneof": "Relation.rel_type", "field": "local_relation", "number": 11, "status": "done", "refused": "", "target": "Read(ReadSource::Local) from the Arrow IPC data plus the DDL or JSON schema (createDataFrame), built by the host before translation (TranslateConf::local_relations); no data and no schema is SCHEMA_REQUIRED_FOR_LOCAL_RELATION; the inline size is bounded by the gRPC message limit", "task": "P3-T012a", "group": 1}, {"oneof": "Relation.rel_type", "field": "sample", "number": 12, "status": "missing", "refused": "", "target": "new Sample (Bernoulli/Poisson with seed, Spark's XORShiftRandom order)", "task": "P3-T012", "group": 3}, {"oneof": "Relation.rel_type", "field": "offset", "number": 13, "status": "done", "refused": "", "target": "Offset", "task": "P2-T091a", "group": 0}, {"oneof": "Relation.rel_type", "field": "deduplicate", "number": 14, "status": "missing", "refused": "", "target": "lowering: Distinct (all columns) or Aggregate with first per non-key column; within_watermark → P4", "task": "P3-T012", "group": 3}, {"oneof": "Relation.rel_type", "field": "range", "number": 15, "status": "done", "refused": "", "target": "Range", "task": "P2-T091a", "group": 0}, - {"oneof": "Relation.rel_type", "field": "subquery_alias", "number": 16, "status": "partial", "refused": "qualifier", "target": "SubqueryAlias with the qualifier", "task": "P3-T012", "group": 1}, + {"oneof": "Relation.rel_type", "field": "subquery_alias", "number": 16, "status": "done", "refused": "", "target": "SubqueryAlias; the qualifier and the alias are the alias parts joined by ALIAS_PART_SEP", "task": "P3-T012", "group": 1}, {"oneof": "Relation.rel_type", "field": "repartition", "number": 17, "status": "missing", "refused": "", "target": "Repartition (round robin / coalesce)", "task": "P3-T012", "group": 3}, {"oneof": "Relation.rel_type", "field": "to_df", "number": 18, "status": "missing", "refused": "", "target": "lowering: Project of Alias (positional)", "task": "P3-T012", "group": 4}, {"oneof": "Relation.rel_type", "field": "with_columns_renamed", "number": 19, "status": "missing", "refused": "", "target": "lowering: Project", "task": "P3-T012", "group": 4}, @@ -33,7 +33,7 @@ {"oneof": "Relation.rel_type", "field": "with_watermark", "number": 33, "status": "out of scope", "refused": "", "target": "P4", "task": "", "group": 0}, {"oneof": "Relation.rel_type", "field": "apply_in_pandas_with_state", "number": 34, "status": "out of scope", "refused": "", "target": "P4", "task": "", "group": 0}, {"oneof": "Relation.rel_type", "field": "html_string", "number": 35, "status": "missing", "refused": "", "target": "as show_string, HTML form", "task": "P3-T012", "group": 4}, - {"oneof": "Relation.rel_type", "field": "cached_local_relation", "number": 36, "status": "missing", "refused": "", "target": "InlineTable from the artifact cache (AddArtifacts cache/, T011); errors NOT_FOUND_CACHED_LOCAL_RELATION", "task": "P3-T083", "group": 7}, + {"oneof": "Relation.rel_type", "field": "cached_local_relation", "number": 36, "status": "partial", "refused": "hashes in the artifact store (AddArtifacts cache/ is not stored yet: every hash is NOT_FOUND_CACHED_LOCAL_RELATION)", "target": "Read(ReadSource::Local) through the same host request as local_relation; the artifact lookup is the host's", "task": "P3-T083", "group": 7}, {"oneof": "Relation.rel_type", "field": "cached_remote_relation", "number": 37, "status": "missing", "refused": "", "target": "the session's cached-DataFrame registry (T083)", "task": "P3-T083", "group": 7}, {"oneof": "Relation.rel_type", "field": "common_inline_user_defined_table_function", "number": 38, "status": "missing", "refused": "", "target": "TableFunction with a Python UDTF (T014)", "task": "P3-T014", "group": 6}, {"oneof": "Relation.rel_type", "field": "as_of_join", "number": 39, "status": "missing", "refused": "", "target": "new AsOfJoin, shared with the ASOF JOIN SQL syntax", "task": "P3-T012", "group": 8}, @@ -61,10 +61,10 @@ {"oneof": "Relation.rel_type", "field": "ml_relation", "number": 300, "status": "out of scope", "refused": "", "target": "machine learning is out of scope", "task": "", "group": 0}, {"oneof": "Relation.rel_type", "field": "extension", "number": 998, "status": "missing", "refused": "", "target": "the krishiv.v1 extension handler; a message no handler claims is refused", "task": "P3-T015", "group": 10}, {"oneof": "Relation.rel_type", "field": "unknown", "number": 999, "status": "out of scope", "refused": "", "target": "Spark falls to case other and answers 56K00 with UNKNOWN(999)", "task": "", "group": 0}, - {"oneof": "Expression.expr_type", "field": "literal", "number": 1, "status": "partial", "refused": "Array, Map, Struct, SpecializedArray, Time", "target": "Literal/TypedLiteral; nested literals after the IR literal supports them; nanosecond timestamps are refused until the types exist", "task": "P3-T012", "group": 1}, + {"oneof": "Expression.expr_type", "field": "literal", "number": 1, "status": "partial", "refused": "Time", "target": "Literal/TypedLiteral; nested literals are Literal::{Array,KeyValueMap,Struct}; nanosecond timestamps are refused until the types exist", "task": "P3-T012", "group": 1}, {"oneof": "Expression.expr_type", "field": "unresolved_attribute", "number": 2, "status": "partial", "refused": "is_metadata_column", "target": "Column; metadata columns (_metadata) → the scan's metadata columns", "task": "P3-T012", "group": 1}, {"oneof": "Expression.expr_type", "field": "unresolved_function", "number": 3, "status": "partial", "refused": "is_user_defined_function (qualified names) and is_internal", "target": "Function with qualifier; internal functions map to Spark's internal names where Krishiv has them, otherwise 56K00-style refusal naming the function", "task": "P3-T012", "group": 1}, - {"oneof": "Expression.expr_type", "field": "expression_string", "number": 4, "status": "missing", "refused": "", "target": "the SQL expression parser + analyzer lowering that select_expr/expr() use (krishiv-session)", "task": "P3-T012", "group": 1}, + {"oneof": "Expression.expr_type", "field": "expression_string", "number": 4, "status": "done", "refused": "", "target": "krishiv_sql::parse_and_lower_named_expression, as Session::select_expr and Session::expr (parser options from TranslateConf)", "task": "P3-T012a", "group": 1}, {"oneof": "Expression.expr_type", "field": "unresolved_star", "number": 5, "status": "done", "refused": "", "target": "Star", "task": "P2-T091a", "group": 0}, {"oneof": "Expression.expr_type", "field": "alias", "number": 6, "status": "partial", "refused": "metadata and several names", "target": "Alias; several names → MultiAlias (exists)", "task": "P3-T012", "group": 1}, {"oneof": "Expression.expr_type", "field": "cast", "number": 7, "status": "done", "refused": "", "target": "Cast", "task": "P2-T091a", "group": 0}, @@ -87,7 +87,7 @@ {"oneof": "Command.command_type", "field": "write_operation", "number": 2, "status": "missing", "refused": "", "target": "DataFrameWriter V1: krishiv-session save/save_as_table/insert_into (dataframe.rs), semantics T023", "task": "P3-T023", "group": 11}, {"oneof": "Command.command_type", "field": "create_dataframe_view", "number": 3, "status": "missing", "refused": "", "target": "create_or_replace_temp_view (exists), global temp views", "task": "P3-T012", "group": 4}, {"oneof": "Command.command_type", "field": "write_operation_v2", "number": 4, "status": "missing", "refused": "", "target": "DataFrameWriterV2 (create/replace/append/overwrite/overwrite_partitions exist in dataframe.rs), T023", "task": "P3-T023", "group": 11}, - {"oneof": "Command.command_type", "field": "sql_command", "number": 5, "status": "missing", "refused": "", "target": "spark.sql() sends this. Run commands at once; for queries return sql_command_result.relation (Spark returns a SQL or LocalRelation relation the client re-sends). Without it no spark.sql test passes", "task": "P3-T012", "group": 1}, + {"oneof": "Command.command_type", "field": "sql_command", "number": 5, "status": "partial", "refused": "input of WithRelations form (group 4; the plain SQL input is supported); named and positional arguments (the analyzer refuses SQL parameters)", "target": "spark.sql() sends this. Run commands at once; for queries return sql_command_result.relation (Spark returns a SQL or LocalRelation relation the client re-sends). Translated here as the SQL relation of a query; the server runs commands (execute.rs)", "task": "P3-T012", "group": 1}, {"oneof": "Command.command_type", "field": "write_stream_operation_start", "number": 6, "status": "out of scope", "refused": "", "target": "P4", "task": "", "group": 0}, {"oneof": "Command.command_type", "field": "streaming_query_command", "number": 7, "status": "out of scope", "refused": "", "target": "P4", "task": "", "group": 0}, {"oneof": "Command.command_type", "field": "get_resources_command", "number": 8, "status": "missing", "refused": "", "target": "an empty resource map (single node, no executors)", "task": "P3-T012", "group": 4}, @@ -156,11 +156,11 @@ {"oneof": "Expression.Literal.literal_type", "field": "calendar_interval", "number": 19, "status": "done", "refused": "", "target": "Literal / TypedLiteral", "task": "P2-T091a", "group": 1}, {"oneof": "Expression.Literal.literal_type", "field": "year_month_interval", "number": 20, "status": "done", "refused": "", "target": "Literal / TypedLiteral", "task": "P2-T091a", "group": 1}, {"oneof": "Expression.Literal.literal_type", "field": "day_time_interval", "number": 21, "status": "done", "refused": "", "target": "Literal / TypedLiteral", "task": "P2-T091a", "group": 1}, - {"oneof": "Expression.Literal.literal_type", "field": "array", "number": 22, "status": "missing", "refused": "", "target": "nested and TIME literals wait for IR values", "task": "P3-T012", "group": 1}, - {"oneof": "Expression.Literal.literal_type", "field": "map", "number": 23, "status": "missing", "refused": "", "target": "nested and TIME literals wait for IR values", "task": "P3-T012", "group": 1}, - {"oneof": "Expression.Literal.literal_type", "field": "struct", "number": 24, "status": "missing", "refused": "", "target": "nested and TIME literals wait for IR values", "task": "P3-T012", "group": 1}, - {"oneof": "Expression.Literal.literal_type", "field": "specialized_array", "number": 25, "status": "missing", "refused": "", "target": "nested and TIME literals wait for IR values", "task": "P3-T012", "group": 1}, - {"oneof": "Expression.Literal.literal_type", "field": "time", "number": 26, "status": "missing", "refused": "", "target": "nested and TIME literals wait for IR values", "task": "P3-T012", "group": 1}, + {"oneof": "Expression.Literal.literal_type", "field": "array", "number": 22, "status": "done", "refused": "", "target": "Literal::{Array,KeyValueMap,Struct}; the type is the root data_type, else the message's own type fields", "task": "P3-T012a", "group": 1}, + {"oneof": "Expression.Literal.literal_type", "field": "map", "number": 23, "status": "done", "refused": "", "target": "Literal::{Array,KeyValueMap,Struct}; the type is the root data_type, else the message's own type fields", "task": "P3-T012a", "group": 1}, + {"oneof": "Expression.Literal.literal_type", "field": "struct", "number": 24, "status": "done", "refused": "", "target": "Literal::{Array,KeyValueMap,Struct}; the type is the root data_type, else the message's own type fields", "task": "P3-T012a", "group": 1}, + {"oneof": "Expression.Literal.literal_type", "field": "specialized_array", "number": 25, "status": "done", "refused": "", "target": "Literal::{Array,KeyValueMap,Struct}; the type is the root data_type, else the message's own type fields", "task": "P3-T012a", "group": 1}, + {"oneof": "Expression.Literal.literal_type", "field": "time", "number": 26, "status": "missing", "refused": "", "target": "TIME literals wait for an IR value", "task": "P3-T012", "group": 1}, {"oneof": "Expression.Literal.literal_type", "field": "timestamp_ntz_nanos", "number": 29, "status": "missing", "refused": "", "target": "no IR type for nanosecond timestamps", "task": "P3-T012", "group": 1}, {"oneof": "Expression.Literal.literal_type", "field": "timestamp_ltz_nanos", "number": 30, "status": "missing", "refused": "", "target": "no IR type for nanosecond timestamps", "task": "P3-T012", "group": 1}, {"oneof": "DataType.kind", "field": "null", "number": 1, "status": "missing", "refused": "", "target": "no IR type", "task": "P3-T012", "group": 1},