huan233usc commented on code in PR #22839: URL: https://github.com/apache/datafusion/pull/22839#discussion_r3382650835
########## datafusion-cli/src/object_storage/stdin.rs: ########## @@ -0,0 +1,230 @@ +// 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::Read; +use std::sync::Arc; + +use datafusion::common::exec_datafusion_err; +use datafusion::config::ConfigFileType; +use datafusion::error::Result; + +use object_store::memory::InMemory; +use object_store::path::Path as ObjectStorePath; +use object_store::{ObjectStore, ObjectStoreExt}; +use url::Url; + +/// 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"; + + /// 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 LOCATIONS: [&'static str; 3] = ["/dev/stdin", "/dev/fd/0", "/proc/self/fd/0"]; + + /// Returns `true` if `location` refers to the process's standard input. + fn is_stdin_location(location: &str) -> bool { + Self::LOCATIONS.contains(&location) + } + + /// 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. + pub(crate) fn rewrite_location( + location: &str, + format: Option<&ConfigFileType>, + ) -> String { + if !Self::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) + } + + /// 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. + pub(crate) async fn object_store(url: &Url) -> Result<Arc<dyn ObjectStore>> { Review Comment: Thanks for the feedback. Yes this is a real correctness bug. I went with the "reuse the buffered store" option. stdin is read once on the first stdin:// table; subsequent stdin:// tables in the same session reuse the already-registered store instead of re-reading (now-EOF) stdin and overwriting it. I added an unit test to cover the scenario below $ printf 'a,b\n1,foo\n2,bar\n' | datafusion-cli -q --command " CREATE EXTERNAL TABLE t STORED AS CSV LOCATION '/dev/stdin' OPTIONS ('format.has_header' 'true'); SELECT count(*) FROM t; -- 2 CREATE EXTERNAL TABLE t2 STORED AS CSV LOCATION '/dev/stdin' OPTIONS ('format.has_header' 'true'); SELECT count(*) FROM t; -- 2 (was 0 before the fix) SELECT count(*) FROM t2; -- 2 " PTAL. Thanks! -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
