diff --git a/js/src/io/geojson.rs b/js/src/io/geojson.rs index ea450473b..1c8e72ee2 100644 --- a/js/src/io/geojson.rs +++ b/js/src/io/geojson.rs @@ -39,8 +39,8 @@ pub fn read_geojson(file: &[u8], batch_size: Option<usize>) -> WasmResult<Table> #[wasm_bindgen(js_name = writeGeoJSON)] pub fn write_geojson(table: Table) -> WasmResult<Vec<u8>> { let (schema, batches) = table.into_inner(); - let mut rust_table = geoarrow::table::Table::try_new(schema, batches)?; + let rust_table = geoarrow::table::Table::try_new(schema, batches)?; let mut output_file: Vec<u8> = vec![]; - _write_geojson(&mut rust_table, &mut output_file)?; + _write_geojson(rust_table, &mut output_file)?; Ok(output_file) } diff --git a/src/io/csv/writer.rs b/src/io/csv/writer.rs index cac0e4859..f376ba035 100644 --- a/src/io/csv/writer.rs +++ b/src/io/csv/writer.rs @@ -5,9 +5,9 @@ use geozero::GeozeroDatasource; use std::io::Write; /// Write a Table to CSV -pub fn write_csv<W: Write>(table: &mut RecordBatchReader, writer: W) -> Result<()> { +pub fn write_csv<W: Write, S: Into<RecordBatchReader>>(stream: S, writer: W) -> Result<()> { let mut csv_writer = CsvWriter::new(writer); - table.process(&mut csv_writer)?; + stream.into().process(&mut csv_writer)?; Ok(()) } @@ -23,7 +23,7 @@ mod test { let mut output_buffer = Vec::new(); let writer = BufWriter::new(&mut output_buffer); - write_csv(&mut table.into(), writer).unwrap(); + write_csv(&table, writer).unwrap(); let output_string = String::from_utf8(output_buffer).unwrap(); println!("{}", output_string); } diff --git a/src/io/flatgeobuf/writer.rs b/src/io/flatgeobuf/writer.rs index f2b92fd6b..9d5f5bae3 100644 --- a/src/io/flatgeobuf/writer.rs +++ b/src/io/flatgeobuf/writer.rs @@ -9,32 +9,33 @@ use crate::schema::GeoSchemaExt; // TODO: always write CRS saved in Table metadata (you can do this by adding an option) /// Write a Table to a FlatGeobuf file. -pub fn write_flatgeobuf<W: Write>( - table: &mut RecordBatchReader, +pub fn write_flatgeobuf<W: Write, S: Into<RecordBatchReader>>( + stream: S, writer: W, name: &str, ) -> Result<()> { - write_flatgeobuf_with_options(table, writer, name, Default::default()) + write_flatgeobuf_with_options(stream, writer, name, Default::default()) } /// Write a Table to a FlatGeobuf file with specific writer options. /// /// Note: this `name` argument is what OGR observes as the layer name of the file. -pub fn write_flatgeobuf_with_options<W: Write>( - table: &mut RecordBatchReader, +pub fn write_flatgeobuf_with_options<W: Write, S: Into<RecordBatchReader>>( + stream: S, writer: W, name: &str, options: FgbWriterOptions, ) -> Result<()> { + let mut stream = stream.into(); let mut fgb = - FgbWriter::create_with_options(name, infer_flatgeobuf_geometry_type(table)?, options)?; - table.process(&mut fgb)?; + FgbWriter::create_with_options(name, infer_flatgeobuf_geometry_type(&stream)?, options)?; + stream.process(&mut fgb)?; fgb.write(writer)?; Ok(()) } -fn infer_flatgeobuf_geometry_type(table: &RecordBatchReader) -> Result<flatgeobuf::GeometryType> { - let schema = table.schema()?; +fn infer_flatgeobuf_geometry_type(stream: &RecordBatchReader) -> Result<flatgeobuf::GeometryType> { + let schema = stream.schema()?; let fields = &schema.fields; let geom_col_idxs = schema.as_ref().geometry_columns(); if geom_col_idxs.len() != 1 { @@ -73,7 +74,7 @@ mod test { let mut output_buffer = Vec::new(); let writer = BufWriter::new(&mut output_buffer); - write_flatgeobuf(&mut table.into(), writer, "name").unwrap(); + write_flatgeobuf(&table, writer, "name").unwrap(); let mut reader = Cursor::new(output_buffer); let new_table = read_flatgeobuf(&mut reader, Default::default()).unwrap(); diff --git a/src/io/geojson/writer.rs b/src/io/geojson/writer.rs index b4591f74c..1b1f122d1 100644 --- a/src/io/geojson/writer.rs +++ b/src/io/geojson/writer.rs @@ -7,9 +7,9 @@ use std::io::Write; /// Write a Table to GeoJSON /// /// Note: Does not reproject to WGS84 for you -pub fn write_geojson<W: Write>(table: &mut RecordBatchReader, writer: W) -> Result<()> { +pub fn write_geojson<W: Write, S: Into<RecordBatchReader>>(stream: S, writer: W) -> Result<()> { let mut geojson = GeoJsonWriter::new(writer); - table.process(&mut geojson)?; + stream.into().process(&mut geojson)?; Ok(()) } @@ -25,7 +25,7 @@ mod test { let mut output_buffer = Vec::new(); let writer = BufWriter::new(&mut output_buffer); - write_geojson(&mut table.into(), writer).unwrap(); + write_geojson(&table, writer).unwrap(); let output_string = String::from_utf8(output_buffer).unwrap(); println!("{}", output_string); } diff --git a/src/io/geojson_lines/writer.rs b/src/io/geojson_lines/writer.rs index 638cc5fe5..2eacc7c4d 100644 --- a/src/io/geojson_lines/writer.rs +++ b/src/io/geojson_lines/writer.rs @@ -6,8 +6,11 @@ use crate::error::Result; use crate::io::stream::RecordBatchReader; /// Write a table to newline-delimited GeoJSON -pub fn write_geojson_lines<W: Write>(table: &mut RecordBatchReader, writer: W) -> Result<()> { +pub fn write_geojson_lines<W: Write, S: Into<RecordBatchReader>>( + stream: S, + writer: W, +) -> Result<()> { let mut geojson_writer = GeoJsonLineWriter::new(writer); - table.process(&mut geojson_writer)?; + stream.into().process(&mut geojson_writer)?; Ok(()) } diff --git a/src/io/ipc/writer.rs b/src/io/ipc/writer.rs index 220a74281..da9bf11ea 100644 --- a/src/io/ipc/writer.rs +++ b/src/io/ipc/writer.rs @@ -6,8 +6,9 @@ use crate::error::{GeoArrowError, Result}; use crate::io::stream::RecordBatchReader; /// Write a Table to an Arrow IPC (Feather v2) file -pub fn write_ipc<W: Write>(table: &mut RecordBatchReader, writer: W) -> Result<()> { - let inner = table +pub fn write_ipc<W: Write, S: Into<RecordBatchReader>>(stream: S, writer: W) -> Result<()> { + let inner = stream + .into() .take() .ok_or(GeoArrowError::General("Closed stream".to_string()))?; @@ -21,8 +22,9 @@ pub fn write_ipc<W: Write>(table: &mut RecordBatchReader, writer: W) -> Result<( } /// Write a Table to an Arrow IPC stream -pub fn write_ipc_stream<W: Write>(table: &mut RecordBatchReader, writer: W) -> Result<()> { - let inner = table +pub fn write_ipc_stream<W: Write, S: Into<RecordBatchReader>>(stream: S, writer: W) -> Result<()> { + let inner = stream + .into() .take() .ok_or(GeoArrowError::General("Closed stream".to_string()))?; diff --git a/src/io/stream.rs b/src/io/stream.rs index f6c689ff4..d15154fb8 100644 --- a/src/io/stream.rs +++ b/src/io/stream.rs @@ -35,6 +35,12 @@ impl From<Table> for RecordBatchReader { } } +impl From<&Table> for RecordBatchReader { + fn from(value: &Table) -> Self { + value.clone().into() + } +} + impl From<Box<dyn _RecordBatchReader>> for RecordBatchReader { fn from(value: Box<dyn _RecordBatchReader>) -> Self { Self(Some(value))