From a0e6d49cd5e99569999ad20c56dff956fc928685 Mon Sep 17 00:00:00 2001 From: Nathan <56370526+nathanb9@users.noreply.github.com> Date: Tue, 16 Jun 2026 15:16:35 -0400 Subject: [PATCH 1/2] Make LogicalPlan::Unnest expression/rebuild contracts consistent (#22783) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Which issue does this PR close? Closes #22769 ## Rationale for this change `LogicalPlan::Unnest` had an inconsistent API contract: `apply_expressions()` exposed `exec_columns` but `with_new_exprs()` rejected them via `assert_no_expressions`. This broke the standard `node.with_new_exprs(node.expressions(), new_inputs)` pattern. ## What changes are included in this PR? - `with_new_exprs` now accepts expressions from `apply_expressions` (extracts `Column` values back out) - `map_expressions` now properly transforms `exec_columns` instead of treating Unnest as expressionless - Removed stale comment in `extract_leaf_expressions` (semantic barrier remains) ## Are these changes tested? Yes — two new unit tests proving both `with_new_exprs(expressions(), inputs)` and `with_new_exprs(vec![], inputs)` work. All existing optimizer and SLT tests pass. ## Are there any user-facing changes? No. --------- Co-authored-by: Andrew Lamb --- datafusion/expr/src/logical_plan/plan.rs | 67 +++++++++++++++++-- datafusion/expr/src/logical_plan/tree_node.rs | 36 +++++++++- .../optimizer/src/extract_leaf_expressions.rs | 18 +++-- 3 files changed, 104 insertions(+), 17 deletions(-) diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index 8dbf41c37f4d1..9ca6941a61ce6 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -1186,12 +1186,20 @@ impl LogicalPlan { options, .. }) => { - self.assert_no_expressions(expr)?; + let exec_columns = if expr.is_empty() { + columns.clone() + } else { + expr.into_iter() + .map(|e| match e { + Expr::Column(c) => Ok(c), + other => internal_err!( + "Expected Expr::Column for Unnest exec_columns, got {other:?}" + ), + }) + .collect::>>()? + }; let input = self.only_input(inputs)?; - // Update schema with unnested column type. - let new_plan = - unnest_with_options(input, columns.clone(), options.clone())?; - Ok(new_plan) + Ok(unnest_with_options(input, exec_columns, options.clone())?) } } } @@ -6602,4 +6610,53 @@ mod tests { Ok(()) } + + #[test] + fn test_unnest_with_new_exprs_accepts_expressions() -> Result<()> { + use crate::LogicalPlanBuilder; + use arrow::datatypes::{DataType, Field, Schema}; + + let schema = Schema::new(vec![ + Field::new("list_col", DataType::new_list(DataType::Int32, true), true), + Field::new("other_col", DataType::Int32, true), + ]); + let plan = table_scan(Some("t"), &schema, None)?.build()?; + let unnest_plan = LogicalPlanBuilder::from(plan) + .unnest_column("list_col")? + .build()?; + + let exprs = unnest_plan.expressions(); + assert!(!exprs.is_empty(), "Unnest should expose exec_columns"); + assert_eq!(exprs.len(), 1); + assert!(matches!(&exprs[0], Expr::Column(c) if c.name == "list_col")); + + let inputs: Vec = + unnest_plan.inputs().into_iter().cloned().collect(); + let rebuilt = unnest_plan.with_new_exprs(exprs, inputs)?; + assert_eq!(rebuilt.schema(), unnest_plan.schema()); + + Ok(()) + } + + #[test] + fn test_unnest_with_new_exprs_empty_preserves_columns() -> Result<()> { + use crate::LogicalPlanBuilder; + use arrow::datatypes::{DataType, Field, Schema}; + + let schema = Schema::new(vec![ + Field::new("list_col", DataType::new_list(DataType::Int32, true), true), + Field::new("other_col", DataType::Int32, true), + ]); + let plan = table_scan(Some("t"), &schema, None)?.build()?; + let unnest_plan = LogicalPlanBuilder::from(plan) + .unnest_column("list_col")? + .build()?; + + let inputs: Vec = + unnest_plan.inputs().into_iter().cloned().collect(); + let rebuilt = unnest_plan.with_new_exprs(vec![], inputs)?; + assert_eq!(rebuilt.schema(), unnest_plan.schema()); + + Ok(()) + } } diff --git a/datafusion/expr/src/logical_plan/tree_node.rs b/datafusion/expr/src/logical_plan/tree_node.rs index cba2dac24b610..c10ac92eef4f5 100644 --- a/datafusion/expr/src/logical_plan/tree_node.rs +++ b/datafusion/expr/src/logical_plan/tree_node.rs @@ -37,13 +37,15 @@ //! * [`LogicalPlan::with_new_exprs`]: Create a new plan with different expressions //! * [`LogicalPlan::expressions`]: Return a copy of the plan's expressions +use std::sync::Arc; + use crate::logical_plan::plan::RangePartitioning; use crate::{ Aggregate, Analyze, CreateMemoryTable, CreateView, DdlStatement, Distinct, DistinctOn, DmlStatement, Execute, Explain, Expr, Extension, Filter, Join, Limit, LogicalPlan, Partitioning, Prepare, Projection, RecursiveQuery, Repartition, Sort, Statement, Subquery, SubqueryAlias, TableScan, Union, Unnest, UserDefinedLogicalNode, - Values, Window, dml::CopyTo, + Values, Window, builder::unnest_with_options, dml::CopyTo, }; use datafusion_common::tree_node::TreeNodeRefContainer; @@ -686,9 +688,39 @@ impl LogicalPlan { _ => Transformed::no(stmt), } .update_data(LogicalPlan::Statement), + LogicalPlan::Unnest(Unnest { + input, + exec_columns, + options, + .. + }) => { + let exprs: Vec = + exec_columns.into_iter().map(Expr::Column).collect(); + exprs.map_elements(f)?.map_data(|mapped_exprs| { + let new_columns = mapped_exprs + .into_iter() + .map(|e| match e { + Expr::Column(c) => Ok(c), + other => internal_err!( + "Expected Expr::Column for Unnest exec_columns, got {other:?}" + ), + }) + .collect::>>()?; + // Rebuild through `unnest_with_options` so the derived + // `list_type_columns`, `struct_type_columns`, + // `dependency_indices`, and `schema` are recomputed from + // the (possibly rewritten) columns rather than carried over + // stale. This keeps `map_expressions` consistent with + // `with_new_exprs`. + unnest_with_options( + Arc::unwrap_or_clone(input), + new_columns, + options, + ) + })? + } // plans without expressions LogicalPlan::EmptyRelation(_) - | LogicalPlan::Unnest(_) | LogicalPlan::RecursiveQuery(_) | LogicalPlan::Subquery(_) | LogicalPlan::SubqueryAlias(_) diff --git a/datafusion/optimizer/src/extract_leaf_expressions.rs b/datafusion/optimizer/src/extract_leaf_expressions.rs index c90f1567fadbb..b855f224c420b 100644 --- a/datafusion/optimizer/src/extract_leaf_expressions.rs +++ b/datafusion/optimizer/src/extract_leaf_expressions.rs @@ -1151,7 +1151,6 @@ fn try_push_into_inputs( // Unnest may output a column with the same name but different value/type // than its input column. Name-based routing cannot distinguish those. - // On top of that Unnest can't go through the `node.with_new_exprs(node.expressions(), new_inputs)` rebuild if matches!(node, LogicalPlan::Unnest(_)) { return Ok(None); } @@ -3046,16 +3045,15 @@ mod tests { Ok(()) } - /// Regression test for the `Assertion failed: expr.is_empty(): Unnest` - /// internal error. + /// Regression test: the optimizer must not push extractions through + /// `Unnest`. /// - /// `try_push_into_inputs` rebuilds the parent node via - /// `node.with_new_exprs(node.expressions(), new_inputs)`. For `Unnest`, - /// `apply_expressions` exposes the `exec_columns` as `Expr::Column`s - /// (so `expressions()` is **non-empty**), but `with_new_exprs` for - /// `Unnest` immediately calls `assert_no_expressions(expr)?` and errors - /// out. The optimizer should treat `Unnest` as a barrier and bail - /// instead of attempting to push through it. + /// `try_push_into_inputs` routes extracted pairs to inputs by column name. + /// `Unnest` can emit an output column with the same name as its input + /// column but a different value/type (the unnested element), so name-based + /// routing cannot tell the two apart. `try_push_into_inputs` therefore + /// treats `Unnest` as a barrier and bails instead of pushing through it + /// (see the `matches!(node, LogicalPlan::Unnest(_))` guard there). #[test] fn test_no_push_through_unnest() -> Result<()> { use arrow::datatypes::{DataType, Field, Schema}; From 96a6096c6f4b924e8cab4bc1629759a948e12939 Mon Sep 17 00:00:00 2001 From: Xin Huang <42597328+huan233usc@users.noreply.github.com> Date: Tue, 16 Jun 2026 12:17:45 -0700 Subject: [PATCH 2/2] feat: support reading from stdin in datafusion-cli (#22839) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Which issue does this PR close? - Closes #9430. ## Rationale for this change Users frequently want to pipe data into the CLI, e.g. `cat data.csv | datafusion-cli`, but pointing `LOCATION` at `/dev/stdin` did not work: - CSV failed with `Illegal seek` (a pipe is not seekable). - Parquet failed with `file size of 0 is less than footer` (a pipe reports size 0). - JSON silently returned 0 rows. This PR makes reading from standard input work for CSV, JSON, and Parquet. ## What changes are included in this PR? stdin is exposed as a `stdin://` object store, dispatched alongside the other schemes (`s3`, `gs`, `http`, ...) in `get_object_store` — conceptually similar to DuckDB's `PipeFileSystem`. - `rewrite_stdin_location` maps the well-known stdin pseudo-paths (`/dev/stdin`, `/dev/fd/0`, `/proc/self/fd/0`) to a canonical `stdin:///stdin.` URL, so they flow through the normal object-store/listing code path. The extension matches the declared `STORED AS` format because the listing layer filters candidate files by extension. - The `stdin://` store reads all of standard input into an in-memory object store. Buffering up front is required because a pipe is not seekable and Parquet stores its metadata at the end of the file. Known scope/limitations (left as potential follow-ups): - Only `CREATE EXTERNAL TABLE` is supported (not dynamic `SELECT * FROM '/dev/stdin'`). - Input is fully buffered in memory, so it must fit in memory. - stdin can only be consumed once per session. - Unix-only (`/dev/stdin` does not exist on Windows); writing to `/dev/stdout` is out of scope. ## Are these changes tested? Yes: - Unit tests in `object_storage.rs` cover `rewrite_stdin_location` and end-to-end reads for CSV, JSON, and Parquet via the in-memory store. - A `#[cfg(unix)]` integration test in `cli_integration.rs` drives the real binary through an actual pipe, exercising the real stdin read. - Manually verified all three formats via real pipes, and confirmed normal local-file reads are unaffected. ## Are there any user-facing changes? Yes — reading from stdin via `LOCATION '/dev/stdin'` is now supported. Documented in `docs/source/user-guide/cli/datasources.md` (new "Reading from standard input" section). No breaking changes. --- datafusion-cli/src/exec.rs | 9 +- datafusion-cli/src/main.rs | 31 +- datafusion-cli/src/object_storage.rs | 6 + datafusion-cli/src/object_storage/stdin.rs | 377 +++++++++++++++++++++ datafusion-cli/tests/cli_integration.rs | 216 ++++++++++++ docs/source/user-guide/cli/datasources.md | 24 ++ 6 files changed, 659 insertions(+), 4 deletions(-) create mode 100644 datafusion-cli/src/object_storage/stdin.rs diff --git a/datafusion-cli/src/exec.rs b/datafusion-cli/src/exec.rs index f7d1541b93bff..f43854821b2d5 100644 --- a/datafusion-cli/src/exec.rs +++ b/datafusion-cli/src/exec.rs @@ -23,7 +23,7 @@ use crate::print_format::PrintFormat; use crate::{ command::{Command, OutputFormat}, helper::CliHelper, - object_storage::get_object_store, + object_storage::{get_object_store, stdin::StdinUtils}, print_options::{MaxRows, PrintOptions}, }; use datafusion::common::instant::Instant; @@ -417,9 +417,14 @@ async fn create_plan( // Note that cmd is a mutable reference so that create_external_table function can remove all // datafusion-cli specific options before passing through to datafusion. Otherwise, datafusion // will raise Configuration errors. - if let LogicalPlan::Ddl(DdlStatement::CreateExternalTable(cmd)) = &plan { + if let LogicalPlan::Ddl(DdlStatement::CreateExternalTable(cmd)) = &mut plan { // To support custom formats, treat error as None let format = config_file_type_from_str(&cmd.file_type); + + // Expose stdin (e.g. `cat data.csv | datafusion-cli`) as a `stdin://` + // object store, registered like any other scheme in `get_object_store`. + cmd.location = StdinUtils::rewrite_location(&cmd.location, format.as_ref()); + register_object_store_and_config_extensions( ctx, &cmd.location, diff --git a/datafusion-cli/src/main.rs b/datafusion-cli/src/main.rs index 0d8ada1367826..4646c5cce9380 100644 --- a/datafusion-cli/src/main.rs +++ b/datafusion-cli/src/main.rs @@ -37,6 +37,7 @@ use datafusion_cli::functions::{ use datafusion_cli::object_storage::instrumented::{ InstrumentedObjectStoreMode, InstrumentedObjectStoreRegistry, }; +use datafusion_cli::object_storage::{StdinCarriesCommands, is_stdin_location}; use datafusion_cli::{ DATAFUSION_CLI_VERSION, exec, pool_type::PoolType, @@ -158,6 +159,23 @@ struct Args { object_store_profiling: InstrumentedObjectStoreMode, } +impl Args { + /// Without -c/-f the CLI enters the REPL, which reads its SQL from + /// stdin — interactively or piped. + fn repl_mode(&self) -> bool { + self.command.is_empty() && self.file.is_empty() + } + + /// Whether the CLI consumes stdin for its own SQL input. This covers the + /// REPL (no -c/-f, reading SQL interactively or piped) as well as an + /// explicit `-f /dev/stdin` (or the other stdin pseudo-paths), where the + /// SQL file *is* stdin. In either case stdin is already spoken for and + /// cannot also back a `LOCATION '/dev/stdin'` table. + fn reads_sql_from_stdin(&self) -> bool { + self.repl_mode() || self.file.iter().any(|f| is_stdin_location(f)) + } +} + #[tokio::main] /// Calls [`main_inner`], then handles printing errors and returning the correct exit code pub async fn main() -> ExitCode { @@ -268,6 +286,7 @@ async fn main_inner() -> Result<()> { instrumented_registry: Arc::clone(&instrumented_registry), }; + let repl_mode = args.repl_mode(); let commands = args.command; let files = args.file; let rc = match args.rc { @@ -285,7 +304,7 @@ async fn main_inner() -> Result<()> { } }; - if commands.is_empty() && files.is_empty() { + if repl_mode { if !rc.is_empty() { exec::exec_from_files(&ctx, rc, &print_options).await?; } @@ -330,8 +349,16 @@ fn get_session_config(args: &Args) -> Result { config_options.format.null = String::from("NULL"); } - let session_config = + let mut session_config = SessionConfig::from(config_options).with_information_schema(true); + + if args.reads_sql_from_stdin() { + // When stdin carries the session's SQL — the REPL (including any rc + // file run before it) or an explicit `-f /dev/stdin` — it cannot also + // serve as a data source for `LOCATION '/dev/stdin'`. + session_config = session_config.with_extension(Arc::new(StdinCarriesCommands)); + } + Ok(session_config) } diff --git a/datafusion-cli/src/object_storage.rs b/datafusion-cli/src/object_storage.rs index 34787838929f1..4293788e0c03a 100644 --- a/datafusion-cli/src/object_storage.rs +++ b/datafusion-cli/src/object_storage.rs @@ -16,6 +16,9 @@ // under the License. pub mod instrumented; +pub(crate) mod stdin; + +pub use stdin::{StdinCarriesCommands, is_stdin_location}; use async_trait::async_trait; use aws_config::BehaviorVersion; @@ -564,6 +567,9 @@ pub(crate) async fn get_object_store( .with_url(url.origin().ascii_serialization()) .build()?, ), + _ if scheme == stdin::StdinUtils::SCHEME => { + stdin::StdinUtils::get_or_create(state, url).await? + } _ => { // For other types, try to get from `object_store_registry`: state diff --git a/datafusion-cli/src/object_storage/stdin.rs b/datafusion-cli/src/object_storage/stdin.rs new file mode 100644 index 0000000000000..602b00a9b90f6 --- /dev/null +++ b/datafusion-cli/src/object_storage/stdin.rs @@ -0,0 +1,377 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Exposes the process's standard input as a `stdin://` object store so that +//! piped data (e.g. `cat data.csv | datafusion-cli`) can be queried via +//! `CREATE EXTERNAL TABLE ... LOCATION '/dev/stdin'`. + +use std::io::{IsTerminal, Read}; +use std::sync::Arc; + +use datafusion::common::exec_datafusion_err; +use datafusion::config::ConfigFileType; +use datafusion::error::Result; +use datafusion::execution::context::SessionState; +use futures::TryStreamExt; + +use object_store::memory::InMemory; +use object_store::path::Path as ObjectStorePath; +use object_store::{ObjectStore, ObjectStoreExt}; +use url::Url; + +/// Marker [`SessionConfig`] extension recording that the session reads its SQL +/// commands from stdin (the interactive or piped REPL). stdin cannot then also +/// serve as a data source: reading it for table data would silently consume +/// the remaining SQL statements. +/// +/// [`SessionConfig`]: datafusion::execution::context::SessionConfig +#[derive(Debug)] +pub struct StdinCarriesCommands; + +/// Filesystem paths that refer to the process's standard input. +/// +/// These are intentionally limited to the well known pseudo-files exposed by +/// the operating system so that ordinary files are never accidentally treated +/// as stdin. +const STDIN_LOCATIONS: [&str; 3] = ["/dev/stdin", "/dev/fd/0", "/proc/self/fd/0"]; + +/// Returns `true` if `path` refers to the process's standard input. +/// +/// Re-exported as [`crate::object_storage::is_stdin_location`] so the CLI entry +/// point can detect when it reads its SQL from stdin via `-f /dev/stdin` and +/// avoid also offering stdin as a `LOCATION '/dev/stdin'` data source. +pub fn is_stdin_location(path: &str) -> bool { + STDIN_LOCATIONS.contains(&path) +} + +/// Utilities for exposing the process's standard input as an object store. +/// +/// stdin is surfaced as a `stdin://` object store and dispatched alongside the +/// other schemes (`s3`, `gs`, `http`, ...) so that reading piped data flows +/// through the normal object-store/listing code path, conceptually similar to +/// DuckDB's `PipeFileSystem`. +pub(crate) struct StdinUtils; + +impl StdinUtils { + /// The URL scheme used to expose stdin as an object store, mirroring how + /// `s3`, `gs`, `http`, etc. are addressed. + pub(crate) const SCHEME: &'static str = "stdin"; + + /// Rewrites the well known stdin pseudo-paths (e.g. `/dev/stdin`) to a + /// canonical `stdin://` URL so that reading from standard input flows + /// through the same object-store/listing code path as any other scheme. + /// Non-stdin locations are returned unchanged. + /// + /// The listing layer filters candidate files by extension, so the canonical + /// object is named with the extension matching the declared `STORED AS` + /// format. The name thereby also records which format stdin was consumed + /// as: a later stdin-backed table declaring a different format resolves to + /// a path the buffered store does not contain and is rejected by + /// [`Self::get_or_create`]. + pub(crate) fn rewrite_location( + location: &str, + format: Option<&ConfigFileType>, + ) -> String { + if !is_stdin_location(location) { + return location.to_string(); + } + + let object_name = match format { + Some(ConfigFileType::CSV) => "stdin.csv", + Some(ConfigFileType::JSON) => "stdin.json", + Some(ConfigFileType::PARQUET) => "stdin.parquet", + _ => "stdin", + }; + format!("{}:///{object_name}", Self::SCHEME) + } + + /// Returns the object store backing the `stdin://` scheme, reading and + /// buffering standard input on first use and reusing that buffer for any + /// subsequent `stdin://` table created in the same session. + /// + /// stdin is a one-shot stream: it can only be read once. The object store + /// registry keys by scheme/authority, so every `stdin://` URL maps to the + /// same store. Without this guard, a second `CREATE EXTERNAL TABLE ... + /// LOCATION '/dev/stdin'` would re-read (now-EOF) stdin, build an empty + /// store, and overwrite the populated one, silently emptying the earlier + /// table. Reusing the already-registered store avoids that. + /// + /// A later stdin-backed table declaring a different `STORED AS` format + /// resolves to an object the store does not contain (the object name + /// records the format stdin was consumed as) and is rejected with a clear + /// error — both reading the buffer as another format and re-reading stdin + /// would be silently wrong. + pub(crate) async fn get_or_create( + state: &SessionState, + url: &Url, + ) -> Result> { + let Ok(existing) = state.runtime_env().object_store_registry.get_store(url) + else { + return Self::object_store(state, url).await; + }; + + let path = ObjectStorePath::from_url_path(url.path())?; + if existing.head(&path).await.is_err() { + let buffered = existing + .list(None) + .try_next() + .await + .ok() + .flatten() + .map(|meta| format!(" as '{}'", meta.location)) + .unwrap_or_default(); + return Err(exec_datafusion_err!( + "stdin was already read{buffered} by an earlier statement; all \ + tables backed by stdin in a session must declare the same \ + STORED AS format" + )); + } + Ok(existing) + } + + /// Builds the object store backing the `stdin://` scheme by reading all of + /// standard input into memory. + /// + /// A pipe (e.g. `cat data.csv | datafusion-cli`) is not seekable and reports + /// a size of `0`, so it cannot be read directly by the file based formats + /// (CSV requires seeking, Parquet needs the footer at the end of the file). + /// Buffering the whole input up front sidesteps these limitations and lets + /// the data be read like any other object, including being scanned more than + /// once. + async fn object_store( + state: &SessionState, + url: &Url, + ) -> Result> { + if state + .config() + .get_extension::() + .is_some() + { + return Err(exec_datafusion_err!( + "stdin is already being read for SQL commands, so it cannot \ + also supply table data; pass the query with -c/--command or \ + -f/--file so that stdin carries the data, e.g. \ + `cat data.csv | datafusion-cli -f query.sql`" + )); + } + if std::io::stdin().is_terminal() { + return Err(exec_datafusion_err!( + "stdin is connected to a terminal, not piped data; pipe the \ + input in, e.g. `cat data.csv | datafusion-cli -f query.sql`" + )); + } + + let mut buffer = Vec::new(); + std::io::stdin() + .lock() + .read_to_end(&mut buffer) + .map_err(|e| exec_datafusion_err!("Failed to read from stdin: {e}"))?; + Self::in_memory_object_store(url, buffer).await + } + + /// Stores `data` at the path referenced by `url` in a fresh [`InMemory`] + /// store. + async fn in_memory_object_store( + url: &Url, + data: Vec, + ) -> Result> { + let store = InMemory::new(); + store + .put(&ObjectStorePath::from_url_path(url.path())?, data.into()) + .await?; + Ok(Arc::new(store)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use datafusion::prelude::{SessionConfig, SessionContext}; + + #[test] + fn rewrites_stdin_locations() { + // stdin pseudo-paths are rewritten to a `stdin://` URL carrying the + // extension that matches the declared format. + assert_eq!( + StdinUtils::rewrite_location("/dev/stdin", Some(&ConfigFileType::CSV)), + "stdin:///stdin.csv" + ); + assert_eq!( + StdinUtils::rewrite_location("/dev/fd/0", Some(&ConfigFileType::JSON)), + "stdin:///stdin.json" + ); + assert_eq!( + StdinUtils::rewrite_location( + "/proc/self/fd/0", + Some(&ConfigFileType::PARQUET) + ), + "stdin:///stdin.parquet" + ); + assert_eq!( + StdinUtils::rewrite_location("/dev/stdin", None), + "stdin:///stdin" + ); + + // Ordinary locations are left untouched. + for location in ["/dev/stdout", "data/stdin.csv", "stdin", "s3://b/f.csv"] { + assert_eq!( + StdinUtils::rewrite_location(location, Some(&ConfigFileType::CSV)), + location + ); + } + } + + /// Buffers `data` into the `stdin://` object store and reads it back through + /// a `CREATE EXTERNAL TABLE`, returning the number of rows in the table. + /// + /// This exercises the full path used for `/dev/stdin` short of the actual + /// stdin read, which cannot be driven from a unit test. + async fn count_stdin_rows( + data: Vec, + stored_as: &str, + format: Option, + options: &str, + ) -> Result { + let location = StdinUtils::rewrite_location("/dev/stdin", format.as_ref()); + let url = Url::parse(&location).unwrap(); + let store = StdinUtils::in_memory_object_store(&url, data).await?; + + let ctx = SessionContext::new(); + ctx.register_object_store(&url, store); + ctx.sql(&format!( + "CREATE EXTERNAL TABLE t STORED AS {stored_as} LOCATION '{location}' {options}" + )) + .await? + .collect() + .await?; + + ctx.sql("SELECT * FROM t").await?.count().await + } + + #[tokio::test] + async fn reuses_buffered_stdin_store() -> Result<()> { + // stdin can only be read once, so a second `stdin://` table must reuse + // the store buffered by the first instead of re-reading (now-empty) + // stdin and overwriting it. + let url = Url::parse("stdin:///stdin.csv").unwrap(); + let store = + StdinUtils::in_memory_object_store(&url, b"a\n1\n2\n".to_vec()).await?; + + let ctx = SessionContext::new(); + ctx.register_object_store(&url, store); + + let reused = StdinUtils::get_or_create(&ctx.state(), &url).await?; + let path = ObjectStorePath::from_url_path(url.path())?; + let bytes = reused.get(&path).await?.bytes().await?; + assert_eq!(bytes.as_ref(), b"a\n1\n2\n"); + Ok(()) + } + + #[tokio::test] + async fn rejects_second_stdin_table_with_different_format() -> Result<()> { + // The buffered object's name records the format stdin was consumed + // as; a later stdin table declaring a different format must fail with + // a clear error rather than a downstream "not found" (or silently + // misreading the bytes as another format). + let csv_url = Url::parse("stdin:///stdin.csv").unwrap(); + let store = + StdinUtils::in_memory_object_store(&csv_url, b"a\n1\n".to_vec()).await?; + + let ctx = SessionContext::new(); + ctx.register_object_store(&csv_url, store); + + let json_url = Url::parse("stdin:///stdin.json").unwrap(); + let err = StdinUtils::get_or_create(&ctx.state(), &json_url) + .await + .unwrap_err() + .to_string(); + assert!( + err.contains("must declare the same STORED AS format") + && err.contains("stdin.csv"), + "unexpected error: {err}" + ); + Ok(()) + } + + #[tokio::test] + async fn errors_when_stdin_carries_commands() { + // Once the REPL owns stdin for SQL commands, building the stdin store + // must fail with a clear error instead of swallowing the remaining + // statements as table data. + let config = SessionConfig::new().with_extension(Arc::new(StdinCarriesCommands)); + let ctx = SessionContext::new_with_config(config); + + let url = Url::parse("stdin:///stdin.csv").unwrap(); + let err = StdinUtils::get_or_create(&ctx.state(), &url) + .await + .unwrap_err(); + assert!( + err.to_string().contains("SQL commands"), + "unexpected error: {err}" + ); + } + + #[tokio::test] + async fn stdin_object_store_reads_csv() -> Result<()> { + let data = b"a,b\n1,foo\n2,bar\n".to_vec(); + let rows = count_stdin_rows( + data, + "CSV", + Some(ConfigFileType::CSV), + "OPTIONS ('format.has_header' 'true')", + ) + .await?; + assert_eq!(rows, 2); + Ok(()) + } + + #[tokio::test] + async fn stdin_object_store_reads_json() -> Result<()> { + let data = b"{\"a\": 1, \"b\": \"foo\"}\n{\"a\": 2, \"b\": \"bar\"}\n".to_vec(); + let rows = count_stdin_rows(data, "JSON", Some(ConfigFileType::JSON), "").await?; + assert_eq!(rows, 2); + Ok(()) + } + + #[tokio::test] + async fn stdin_object_store_reads_parquet() -> Result<()> { + use datafusion::arrow::array::Int32Array; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::arrow::record_batch::RecordBatch; + use parquet::arrow::ArrowWriter; + + // Parquet requires random access to the footer, which a real pipe cannot + // provide; the in-memory buffer makes this work. + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + ) + .unwrap(); + + let mut data = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut data, schema, None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let rows = + count_stdin_rows(data, "PARQUET", Some(ConfigFileType::PARQUET), "").await?; + assert_eq!(rows, 3); + Ok(()) + } +} diff --git a/datafusion-cli/tests/cli_integration.rs b/datafusion-cli/tests/cli_integration.rs index 4849ac9e9a5e2..4dc244445a2eb 100644 --- a/datafusion-cli/tests/cli_integration.rs +++ b/datafusion-cli/tests/cli_integration.rs @@ -173,6 +173,222 @@ fn cli_quick_test<'a>( assert_cmd_snapshot!(cmd); } +/// Read data piped into the CLI via the `/dev/stdin` pseudo-path. +/// +/// Unix-only: `/dev/stdin` does not exist on Windows. This drives the real +/// binary through an actual pipe, exercising the stdin read that the in-process +/// unit tests cannot. +#[cfg(unix)] +#[test] +fn test_cli_read_from_stdin() { + let stdout = run_cli_with_stdin( + "CREATE EXTERNAL TABLE t STORED AS CSV LOCATION '/dev/stdin' \ + OPTIONS ('format.has_header' 'true'); \ + SELECT b, count(*) AS c FROM t GROUP BY b ORDER BY b;", + b"a,b\n1,foo\n2,bar\n3,foo\n", + ); + + assert!( + stdout.contains("| foo | 2 |") && stdout.contains("| bar | 1 |"), + "unexpected output:\n{stdout}" + ); +} + +/// stdin is a one-shot stream, so a second `/dev/stdin` table in the same +/// session must reuse the buffered input rather than re-reading (now-empty) +/// stdin and silently emptying the first table. +#[cfg(unix)] +#[test] +fn test_cli_read_from_stdin_twice_reuses_buffer() { + let stdout = run_cli_with_stdin( + "CREATE EXTERNAL TABLE t STORED AS CSV LOCATION '/dev/stdin' \ + OPTIONS ('format.has_header' 'true'); \ + CREATE EXTERNAL TABLE t2 STORED AS CSV LOCATION '/dev/stdin' \ + OPTIONS ('format.has_header' 'true'); \ + SELECT count(*) AS t_count FROM t; \ + SELECT count(*) AS t2_count FROM t2;", + b"a,b\n1,foo\n2,bar\n", + ); + + // Both tables must still see the two buffered rows. + let counts: Vec<&str> = stdout + .lines() + .filter(|line| line.trim_start().starts_with("| 2 ")) + .collect(); + assert_eq!( + counts.len(), + 2, + "expected both stdin tables to report 2 rows, got:\n{stdout}" + ); +} + +/// A later `/dev/stdin` table declaring a different `STORED AS` format must be +/// rejected with a clear error: stdin is one-shot, its bytes were already +/// buffered under the first table's format, and silently reading them as +/// another format would be wrong. +#[cfg(unix)] +#[test] +fn test_cli_read_from_stdin_mixed_formats_rejected() { + use std::io::Write; + use std::process::Stdio; + + let mut child = cli() + .args([ + "-q", + "--command", + "CREATE EXTERNAL TABLE t STORED AS CSV LOCATION '/dev/stdin' \ + OPTIONS ('format.has_header' 'true'); \ + CREATE EXTERNAL TABLE t2 STORED AS JSON LOCATION '/dev/stdin';", + ]) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .expect("failed to spawn datafusion-cli"); + + child + .stdin + .take() + .unwrap() + .write_all(b"a,b\n1,foo\n2,bar\n") + .unwrap(); + + let output = child.wait_with_output().unwrap(); + // Fatal errors in `--command` mode are reported on stdout. + let stdout = String::from_utf8_lossy(&output.stdout); + + assert!( + !output.status.success(), + "expected the mismatched format to fail, stdout:\n{stdout}" + ); + assert!( + stdout.contains("must declare the same STORED AS format"), + "expected a clear mismatch error, got:\n{stdout}" + ); +} + +/// When the SQL itself arrives on stdin (the piped REPL, e.g. `cat script.sql +/// | datafusion-cli`), stdin cannot double as a data source: the statement +/// must fail with a clear error instead of silently consuming the rest of the +/// script as table data, and the remaining statements must still run. +#[cfg(unix)] +#[test] +fn test_cli_stdin_location_rejected_when_sql_comes_from_stdin() { + use std::io::Write; + use std::process::Stdio; + + let mut child = cli() + .arg("-q") + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .expect("failed to spawn datafusion-cli"); + + child + .stdin + .take() + .unwrap() + .write_all( + b"CREATE EXTERNAL TABLE t STORED AS CSV LOCATION '/dev/stdin';\n\ + SELECT 123 + 456;\n", + ) + .unwrap(); + + let output = child.wait_with_output().unwrap(); + let stdout = String::from_utf8_lossy(&output.stdout); + let stderr = String::from_utf8_lossy(&output.stderr); + + assert!( + stderr.contains("SQL commands"), + "expected a clear error about stdin carrying SQL.\nstdout:\n{stdout}\nstderr:\n{stderr}" + ); + // The statement after the failed CREATE must still execute rather than + // being consumed as table data. + assert!( + stdout.contains("579"), + "expected the following statement to still run.\nstdout:\n{stdout}\nstderr:\n{stderr}" + ); +} + +/// `-f /dev/stdin` reads the SQL script from stdin, exactly like the piped +/// REPL, so stdin still cannot double as a `LOCATION '/dev/stdin'` data source. +/// The offending statement must fail with the same clear error, and later +/// statements in the script must still run. +/// +/// `/dev/stdin` only passes the `-f` file check when stdin is a redirected +/// regular file (a pipe is not `is_file()`), so the binary is driven with a +/// temp script file as its stdin rather than a pipe. +#[cfg(unix)] +#[test] +fn test_cli_dash_f_stdin_location_rejected() { + use std::process::Stdio; + + let script = env::temp_dir().join(format!( + "datafusion_cli_dash_f_stdin_{}.sql", + std::process::id() + )); + fs::write( + &script, + b"CREATE EXTERNAL TABLE t STORED AS CSV LOCATION '/dev/stdin';\n\ + SELECT 123 + 456;\n", + ) + .unwrap(); + let stdin = fs::File::open(&script).unwrap(); + + let output = cli() + .args(["-q", "-f", "/dev/stdin"]) + .stdin(Stdio::from(stdin)) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .output() + .expect("failed to spawn datafusion-cli"); + + let _ = fs::remove_file(&script); + + let stdout = String::from_utf8_lossy(&output.stdout); + let stderr = String::from_utf8_lossy(&output.stderr); + + assert!( + stderr.contains("SQL commands"), + "expected a clear error about stdin carrying SQL.\nstdout:\n{stdout}\nstderr:\n{stderr}" + ); + // The statement after the failed CREATE must still execute rather than + // being consumed as table data. + assert!( + stdout.contains("579"), + "expected the following statement to still run.\nstdout:\n{stdout}\nstderr:\n{stderr}" + ); +} + +/// Spawns the real `datafusion-cli` binary, pipes `stdin` into it, and returns +/// its stdout after asserting a successful exit. +#[cfg(unix)] +fn run_cli_with_stdin(command: &str, stdin: &[u8]) -> String { + use std::io::Write; + use std::process::Stdio; + + let mut child = cli() + .args(["-q", "--command", command]) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .expect("failed to spawn datafusion-cli"); + + child.stdin.take().unwrap().write_all(stdin).unwrap(); + + let output = child.wait_with_output().unwrap(); + let stdout = String::from_utf8_lossy(&output.stdout).into_owned(); + let stderr = String::from_utf8_lossy(&output.stderr); + + assert!( + output.status.success(), + "datafusion-cli failed.\nstdout:\n{stdout}\nstderr:\n{stderr}" + ); + stdout +} + #[test] fn cli_explain_environment_overrides() { let mut settings = make_settings(); diff --git a/docs/source/user-guide/cli/datasources.md b/docs/source/user-guide/cli/datasources.md index 6b1a4887a8a0f..59a6b0aa43284 100644 --- a/docs/source/user-guide/cli/datasources.md +++ b/docs/source/user-guide/cli/datasources.md @@ -132,6 +132,30 @@ select count(*) from hits; 1 row in set. Query took 0.344 seconds. ``` +## Reading from standard input + +On Unix-like systems you can pipe data into the CLI and query it by pointing the +`LOCATION` at the `/dev/stdin` pseudo-file: + +```console +$ cat hits.csv | datafusion-cli -c " +CREATE EXTERNAL TABLE hits STORED AS CSV LOCATION '/dev/stdin' OPTIONS ('format.has_header' 'true'); +SELECT count(*) FROM hits;" +``` + +This works for CSV, JSON, and Parquet. Because standard input is not seekable +(and Parquet stores its metadata at the end of the file), the CLI buffers the +entire input into memory before querying it, so the data must fit in memory. +Standard input is read only once: the buffered contents are reused for any +further tables backed by `/dev/stdin` in the same session. Those tables must +declare the same `STORED AS` format as the first one; a differing format is +rejected with an error. + +The SQL must be passed with `-c`/`--command` or `-f`/`--file` so that standard +input is free to carry the data. In the interactive shell (and when SQL is +piped to the CLI without `-c`/`-f`) standard input carries the SQL itself, and +`LOCATION '/dev/stdin'` returns an error. + **Why Wildcards Are Not Supported** Although wildcards (e.g., _.parquet or \*\*/_.parquet) may work for local