acc3ca9de7
<!-- if this PR closes one or more issues, you can automatically link the PR with them by using one of the [*linking keywords*](https://docs.github.com/en/issues/tracking-your-work-with-issues/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword), e.g. - this PR should close #xxxx - fixes #xxxx you can also mention related issues, PRs or discussions! --> # Description <!-- Thank you for improving Nushell. Please, check our [contributing guide](../CONTRIBUTING.md) and talk to the core team before making major changes. Description of your pull request goes here. **Provide examples and/or screenshots** if your changes affect the user experience. --> The error message when using `dfr open --type` shows an outdated list of supported formats. # User-Facing Changes <!-- List of all changes that impact the user experience here. This helps us keep track of breaking changes. --> User is now informed that jsonl and avro formats are supported. # Tests + Formatting <!-- Don't forget to add tests that cover your changes. Make sure you've run and fixed any issues with these commands: - `cargo fmt --all -- --check` to check standard code formatting (`cargo fmt --all` applies these changes) - `cargo clippy --workspace -- -D warnings -D clippy::unwrap_used` to check that you're using the standard code style - `cargo test --workspace` to check that all tests pass (on Windows make sure to [enable developer mode](https://learn.microsoft.com/en-us/windows/apps/get-started/developer-mode-features-and-debugging)) - `cargo run -- -c "use std testing; testing run-tests --path crates/nu-std"` to run the tests for the standard library > **Note** > from `nushell` you can also use the `toolkit` as follows > ```bash > use toolkit.nu # or use an `env_change` hook to activate it automatically > toolkit check pr > ``` --> Done. # After Submitting <!-- If your PR had any user-facing changes, update [the documentation](https://github.com/nushell/nushell.github.io) after the PR is merged, if necessary. This will help us keep the docs up to date. --> No doc changes.
518 lines
16 KiB
Rust
518 lines
16 KiB
Rust
use crate::dataframe::values::{NuDataFrame, NuLazyFrame, NuSchema};
|
|
use nu_engine::command_prelude::*;
|
|
|
|
use polars::prelude::{
|
|
CsvEncoding, CsvReader, IpcReader, JsonFormat, JsonReader, LazyCsvReader, LazyFileListReader,
|
|
LazyFrame, ParallelStrategy, ParquetReader, ScanArgsIpc, ScanArgsParquet, SerReader,
|
|
};
|
|
use polars_io::avro::AvroReader;
|
|
use std::{fs::File, io::BufReader, path::PathBuf};
|
|
|
|
#[derive(Clone)]
|
|
pub struct OpenDataFrame;
|
|
|
|
impl Command for OpenDataFrame {
|
|
fn name(&self) -> &str {
|
|
"dfr open"
|
|
}
|
|
|
|
fn usage(&self) -> &str {
|
|
"Opens CSV, JSON, JSON lines, arrow, avro, or parquet file to create dataframe."
|
|
}
|
|
|
|
fn signature(&self) -> Signature {
|
|
Signature::build(self.name())
|
|
.required(
|
|
"file",
|
|
SyntaxShape::Filepath,
|
|
"file path to load values from",
|
|
)
|
|
.switch("lazy", "creates a lazy dataframe", Some('l'))
|
|
.named(
|
|
"type",
|
|
SyntaxShape::String,
|
|
"File type: csv, tsv, json, parquet, arrow, avro. If omitted, derive from file extension",
|
|
Some('t'),
|
|
)
|
|
.named(
|
|
"delimiter",
|
|
SyntaxShape::String,
|
|
"file delimiter character. CSV file",
|
|
Some('d'),
|
|
)
|
|
.switch(
|
|
"no-header",
|
|
"Indicates if file doesn't have header. CSV file",
|
|
None,
|
|
)
|
|
.named(
|
|
"infer-schema",
|
|
SyntaxShape::Number,
|
|
"Number of rows to infer the schema of the file. CSV file",
|
|
None,
|
|
)
|
|
.named(
|
|
"skip-rows",
|
|
SyntaxShape::Number,
|
|
"Number of rows to skip from file. CSV file",
|
|
None,
|
|
)
|
|
.named(
|
|
"columns",
|
|
SyntaxShape::List(Box::new(SyntaxShape::String)),
|
|
"Columns to be selected from csv file. CSV and Parquet file",
|
|
None,
|
|
)
|
|
.named(
|
|
"schema",
|
|
SyntaxShape::Record(vec![]),
|
|
r#"Polars Schema in format [{name: str}]. CSV, JSON, and JSONL files"#,
|
|
Some('s')
|
|
)
|
|
.input_output_type(Type::Any, Type::Custom("dataframe".into()))
|
|
.category(Category::Custom("dataframe".into()))
|
|
}
|
|
|
|
fn examples(&self) -> Vec<Example> {
|
|
vec![Example {
|
|
description: "Takes a file name and creates a dataframe",
|
|
example: "dfr open test.csv",
|
|
result: None,
|
|
}]
|
|
}
|
|
|
|
fn run(
|
|
&self,
|
|
engine_state: &EngineState,
|
|
stack: &mut Stack,
|
|
call: &Call,
|
|
_input: PipelineData,
|
|
) -> Result<PipelineData, ShellError> {
|
|
command(engine_state, stack, call)
|
|
}
|
|
}
|
|
|
|
fn command(
|
|
engine_state: &EngineState,
|
|
stack: &mut Stack,
|
|
call: &Call,
|
|
) -> Result<PipelineData, ShellError> {
|
|
let file: Spanned<PathBuf> = call.req(engine_state, stack, 0)?;
|
|
|
|
let type_option: Option<Spanned<String>> = call.get_flag(engine_state, stack, "type")?;
|
|
|
|
let type_id = match &type_option {
|
|
Some(ref t) => Some((t.item.to_owned(), "Invalid type", t.span)),
|
|
None => file.item.extension().map(|e| {
|
|
(
|
|
e.to_string_lossy().into_owned(),
|
|
"Invalid extension",
|
|
file.span,
|
|
)
|
|
}),
|
|
};
|
|
|
|
match type_id {
|
|
Some((e, msg, blamed)) => match e.as_str() {
|
|
"csv" | "tsv" => from_csv(engine_state, stack, call),
|
|
"parquet" | "parq" => from_parquet(engine_state, stack, call),
|
|
"ipc" | "arrow" => from_ipc(engine_state, stack, call),
|
|
"json" => from_json(engine_state, stack, call),
|
|
"jsonl" => from_jsonl(engine_state, stack, call),
|
|
"avro" => from_avro(engine_state, stack, call),
|
|
_ => Err(ShellError::FileNotFoundCustom {
|
|
msg: format!(
|
|
"{msg}. Supported values: csv, tsv, parquet, ipc, arrow, json, jsonl, avro"
|
|
),
|
|
span: blamed,
|
|
}),
|
|
},
|
|
None => Err(ShellError::FileNotFoundCustom {
|
|
msg: "File without extension".into(),
|
|
span: file.span,
|
|
}),
|
|
}
|
|
.map(|value| PipelineData::Value(value, None))
|
|
}
|
|
|
|
fn from_parquet(
|
|
engine_state: &EngineState,
|
|
stack: &mut Stack,
|
|
call: &Call,
|
|
) -> Result<Value, ShellError> {
|
|
if call.has_flag(engine_state, stack, "lazy")? {
|
|
let file: String = call.req(engine_state, stack, 0)?;
|
|
let args = ScanArgsParquet {
|
|
n_rows: None,
|
|
cache: true,
|
|
parallel: ParallelStrategy::Auto,
|
|
rechunk: false,
|
|
row_index: None,
|
|
low_memory: false,
|
|
cloud_options: None,
|
|
use_statistics: false,
|
|
hive_partitioning: false,
|
|
};
|
|
|
|
let df: NuLazyFrame = LazyFrame::scan_parquet(file, args)
|
|
.map_err(|e| ShellError::GenericError {
|
|
error: "Parquet reader error".into(),
|
|
msg: format!("{e:?}"),
|
|
span: Some(call.head),
|
|
help: None,
|
|
inner: vec![],
|
|
})?
|
|
.into();
|
|
|
|
df.into_value(call.head)
|
|
} else {
|
|
let file: Spanned<PathBuf> = call.req(engine_state, stack, 0)?;
|
|
let columns: Option<Vec<String>> = call.get_flag(engine_state, stack, "columns")?;
|
|
|
|
let r = File::open(&file.item).map_err(|e| ShellError::GenericError {
|
|
error: "Error opening file".into(),
|
|
msg: e.to_string(),
|
|
span: Some(file.span),
|
|
help: None,
|
|
inner: vec![],
|
|
})?;
|
|
let reader = ParquetReader::new(r);
|
|
|
|
let reader = match columns {
|
|
None => reader,
|
|
Some(columns) => reader.with_columns(Some(columns)),
|
|
};
|
|
|
|
let df: NuDataFrame = reader
|
|
.finish()
|
|
.map_err(|e| ShellError::GenericError {
|
|
error: "Parquet reader error".into(),
|
|
msg: format!("{e:?}"),
|
|
span: Some(call.head),
|
|
help: None,
|
|
inner: vec![],
|
|
})?
|
|
.into();
|
|
|
|
Ok(df.into_value(call.head))
|
|
}
|
|
}
|
|
|
|
fn from_avro(
|
|
engine_state: &EngineState,
|
|
stack: &mut Stack,
|
|
call: &Call,
|
|
) -> Result<Value, ShellError> {
|
|
let file: Spanned<PathBuf> = call.req(engine_state, stack, 0)?;
|
|
let columns: Option<Vec<String>> = call.get_flag(engine_state, stack, "columns")?;
|
|
|
|
let r = File::open(&file.item).map_err(|e| ShellError::GenericError {
|
|
error: "Error opening file".into(),
|
|
msg: e.to_string(),
|
|
span: Some(file.span),
|
|
help: None,
|
|
inner: vec![],
|
|
})?;
|
|
let reader = AvroReader::new(r);
|
|
|
|
let reader = match columns {
|
|
None => reader,
|
|
Some(columns) => reader.with_columns(Some(columns)),
|
|
};
|
|
|
|
let df: NuDataFrame = reader
|
|
.finish()
|
|
.map_err(|e| ShellError::GenericError {
|
|
error: "Avro reader error".into(),
|
|
msg: format!("{e:?}"),
|
|
span: Some(call.head),
|
|
help: None,
|
|
inner: vec![],
|
|
})?
|
|
.into();
|
|
|
|
Ok(df.into_value(call.head))
|
|
}
|
|
|
|
fn from_ipc(
|
|
engine_state: &EngineState,
|
|
stack: &mut Stack,
|
|
call: &Call,
|
|
) -> Result<Value, ShellError> {
|
|
if call.has_flag(engine_state, stack, "lazy")? {
|
|
let file: String = call.req(engine_state, stack, 0)?;
|
|
let args = ScanArgsIpc {
|
|
n_rows: None,
|
|
cache: true,
|
|
rechunk: false,
|
|
row_index: None,
|
|
memmap: true,
|
|
};
|
|
|
|
let df: NuLazyFrame = LazyFrame::scan_ipc(file, args)
|
|
.map_err(|e| ShellError::GenericError {
|
|
error: "IPC reader error".into(),
|
|
msg: format!("{e:?}"),
|
|
span: Some(call.head),
|
|
help: None,
|
|
inner: vec![],
|
|
})?
|
|
.into();
|
|
|
|
df.into_value(call.head)
|
|
} else {
|
|
let file: Spanned<PathBuf> = call.req(engine_state, stack, 0)?;
|
|
let columns: Option<Vec<String>> = call.get_flag(engine_state, stack, "columns")?;
|
|
|
|
let r = File::open(&file.item).map_err(|e| ShellError::GenericError {
|
|
error: "Error opening file".into(),
|
|
msg: e.to_string(),
|
|
span: Some(file.span),
|
|
help: None,
|
|
inner: vec![],
|
|
})?;
|
|
let reader = IpcReader::new(r);
|
|
|
|
let reader = match columns {
|
|
None => reader,
|
|
Some(columns) => reader.with_columns(Some(columns)),
|
|
};
|
|
|
|
let df: NuDataFrame = reader
|
|
.finish()
|
|
.map_err(|e| ShellError::GenericError {
|
|
error: "IPC reader error".into(),
|
|
msg: format!("{e:?}"),
|
|
span: Some(call.head),
|
|
help: None,
|
|
inner: vec![],
|
|
})?
|
|
.into();
|
|
|
|
Ok(df.into_value(call.head))
|
|
}
|
|
}
|
|
|
|
fn from_json(
|
|
engine_state: &EngineState,
|
|
stack: &mut Stack,
|
|
call: &Call,
|
|
) -> Result<Value, ShellError> {
|
|
let file: Spanned<PathBuf> = call.req(engine_state, stack, 0)?;
|
|
let file = File::open(&file.item).map_err(|e| ShellError::GenericError {
|
|
error: "Error opening file".into(),
|
|
msg: e.to_string(),
|
|
span: Some(file.span),
|
|
help: None,
|
|
inner: vec![],
|
|
})?;
|
|
let maybe_schema = call
|
|
.get_flag(engine_state, stack, "schema")?
|
|
.map(|schema| NuSchema::try_from(&schema))
|
|
.transpose()?;
|
|
|
|
let buf_reader = BufReader::new(file);
|
|
let reader = JsonReader::new(buf_reader);
|
|
|
|
let reader = match maybe_schema {
|
|
Some(schema) => reader.with_schema(schema.into()),
|
|
None => reader,
|
|
};
|
|
|
|
let df: NuDataFrame = reader
|
|
.finish()
|
|
.map_err(|e| ShellError::GenericError {
|
|
error: "Json reader error".into(),
|
|
msg: format!("{e:?}"),
|
|
span: Some(call.head),
|
|
help: None,
|
|
inner: vec![],
|
|
})?
|
|
.into();
|
|
|
|
Ok(df.into_value(call.head))
|
|
}
|
|
|
|
fn from_jsonl(
|
|
engine_state: &EngineState,
|
|
stack: &mut Stack,
|
|
call: &Call,
|
|
) -> Result<Value, ShellError> {
|
|
let infer_schema: Option<usize> = call.get_flag(engine_state, stack, "infer-schema")?;
|
|
let maybe_schema = call
|
|
.get_flag(engine_state, stack, "schema")?
|
|
.map(|schema| NuSchema::try_from(&schema))
|
|
.transpose()?;
|
|
let file: Spanned<PathBuf> = call.req(engine_state, stack, 0)?;
|
|
let file = File::open(&file.item).map_err(|e| ShellError::GenericError {
|
|
error: "Error opening file".into(),
|
|
msg: e.to_string(),
|
|
span: Some(file.span),
|
|
help: None,
|
|
inner: vec![],
|
|
})?;
|
|
|
|
let buf_reader = BufReader::new(file);
|
|
let reader = JsonReader::new(buf_reader)
|
|
.with_json_format(JsonFormat::JsonLines)
|
|
.infer_schema_len(infer_schema);
|
|
|
|
let reader = match maybe_schema {
|
|
Some(schema) => reader.with_schema(schema.into()),
|
|
None => reader,
|
|
};
|
|
|
|
let df: NuDataFrame = reader
|
|
.finish()
|
|
.map_err(|e| ShellError::GenericError {
|
|
error: "Json lines reader error".into(),
|
|
msg: format!("{e:?}"),
|
|
span: Some(call.head),
|
|
help: None,
|
|
inner: vec![],
|
|
})?
|
|
.into();
|
|
|
|
Ok(df.into_value(call.head))
|
|
}
|
|
|
|
fn from_csv(
|
|
engine_state: &EngineState,
|
|
stack: &mut Stack,
|
|
call: &Call,
|
|
) -> Result<Value, ShellError> {
|
|
let delimiter: Option<Spanned<String>> = call.get_flag(engine_state, stack, "delimiter")?;
|
|
let no_header: bool = call.has_flag(engine_state, stack, "no-header")?;
|
|
let infer_schema: Option<usize> = call.get_flag(engine_state, stack, "infer-schema")?;
|
|
let skip_rows: Option<usize> = call.get_flag(engine_state, stack, "skip-rows")?;
|
|
let columns: Option<Vec<String>> = call.get_flag(engine_state, stack, "columns")?;
|
|
|
|
let maybe_schema = call
|
|
.get_flag(engine_state, stack, "schema")?
|
|
.map(|schema| NuSchema::try_from(&schema))
|
|
.transpose()?;
|
|
|
|
if call.has_flag(engine_state, stack, "lazy")? {
|
|
let file: String = call.req(engine_state, stack, 0)?;
|
|
let csv_reader = LazyCsvReader::new(file);
|
|
|
|
let csv_reader = match delimiter {
|
|
None => csv_reader,
|
|
Some(d) => {
|
|
if d.item.len() != 1 {
|
|
return Err(ShellError::GenericError {
|
|
error: "Incorrect delimiter".into(),
|
|
msg: "Delimiter has to be one character".into(),
|
|
span: Some(d.span),
|
|
help: None,
|
|
inner: vec![],
|
|
});
|
|
} else {
|
|
let delimiter = match d.item.chars().next() {
|
|
Some(d) => d as u8,
|
|
None => unreachable!(),
|
|
};
|
|
csv_reader.with_separator(delimiter)
|
|
}
|
|
}
|
|
};
|
|
|
|
let csv_reader = csv_reader.has_header(!no_header);
|
|
|
|
let csv_reader = match maybe_schema {
|
|
Some(schema) => csv_reader.with_schema(Some(schema.into())),
|
|
None => csv_reader,
|
|
};
|
|
|
|
let csv_reader = match infer_schema {
|
|
None => csv_reader,
|
|
Some(r) => csv_reader.with_infer_schema_length(Some(r)),
|
|
};
|
|
|
|
let csv_reader = match skip_rows {
|
|
None => csv_reader,
|
|
Some(r) => csv_reader.with_skip_rows(r),
|
|
};
|
|
|
|
let df: NuLazyFrame = csv_reader
|
|
.finish()
|
|
.map_err(|e| ShellError::GenericError {
|
|
error: "Parquet reader error".into(),
|
|
msg: format!("{e:?}"),
|
|
span: Some(call.head),
|
|
help: None,
|
|
inner: vec![],
|
|
})?
|
|
.into();
|
|
|
|
df.into_value(call.head)
|
|
} else {
|
|
let file: Spanned<PathBuf> = call.req(engine_state, stack, 0)?;
|
|
let csv_reader = CsvReader::from_path(&file.item)
|
|
.map_err(|e| ShellError::GenericError {
|
|
error: "Error creating CSV reader".into(),
|
|
msg: e.to_string(),
|
|
span: Some(file.span),
|
|
help: None,
|
|
inner: vec![],
|
|
})?
|
|
.with_encoding(CsvEncoding::LossyUtf8);
|
|
|
|
let csv_reader = match delimiter {
|
|
None => csv_reader,
|
|
Some(d) => {
|
|
if d.item.len() != 1 {
|
|
return Err(ShellError::GenericError {
|
|
error: "Incorrect delimiter".into(),
|
|
msg: "Delimiter has to be one character".into(),
|
|
span: Some(d.span),
|
|
help: None,
|
|
inner: vec![],
|
|
});
|
|
} else {
|
|
let delimiter = match d.item.chars().next() {
|
|
Some(d) => d as u8,
|
|
None => unreachable!(),
|
|
};
|
|
csv_reader.with_separator(delimiter)
|
|
}
|
|
}
|
|
};
|
|
|
|
let csv_reader = csv_reader.has_header(!no_header);
|
|
|
|
let csv_reader = match maybe_schema {
|
|
Some(schema) => csv_reader.with_schema(Some(schema.into())),
|
|
None => csv_reader,
|
|
};
|
|
|
|
let csv_reader = match infer_schema {
|
|
None => csv_reader,
|
|
Some(r) => csv_reader.infer_schema(Some(r)),
|
|
};
|
|
|
|
let csv_reader = match skip_rows {
|
|
None => csv_reader,
|
|
Some(r) => csv_reader.with_skip_rows(r),
|
|
};
|
|
|
|
let csv_reader = match columns {
|
|
None => csv_reader,
|
|
Some(columns) => csv_reader.with_columns(Some(columns)),
|
|
};
|
|
|
|
let df: NuDataFrame = csv_reader
|
|
.finish()
|
|
.map_err(|e| ShellError::GenericError {
|
|
error: "Parquet reader error".into(),
|
|
msg: format!("{e:?}"),
|
|
span: Some(call.head),
|
|
help: None,
|
|
inner: vec![],
|
|
})?
|
|
.into();
|
|
|
|
Ok(df.into_value(call.head))
|
|
}
|
|
}
|