You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@sedona.apache.org by ji...@apache.org on 2026/07/27 05:16:38 UTC
(sedona-db) branch main updated: feat(python/sedonadb): add DataFrame.to_csv and to_json (#1044)
This is an automated email from the ASF dual-hosted git repository.
jiayuasu pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/sedona-db.git
The following commit(s) were added to refs/heads/main by this push:
new 2e650e02a feat(python/sedonadb): add DataFrame.to_csv and to_json (#1044)
2e650e02a is described below
commit 2e650e02a217cf6b91d50227e5062af72321f549
Author: Jia Yu <ji...@wherobots.com>
AuthorDate: Sun Jul 26 22:16:32 2026 -0700
feat(python/sedonadb): add DataFrame.to_csv and to_json (#1044)
---
python/sedonadb/python/sedonadb/dataframe.py | 60 +++++++++++
python/sedonadb/src/dataframe.rs | 33 ++++++
python/sedonadb/tests/io/test_write_csv_json.py | 128 ++++++++++++++++++++++++
rust/sedona-schema/src/schema.rs | 88 +++++++++++++++-
rust/sedona/src/context.rs | 69 ++++++++++++-
5 files changed, 376 insertions(+), 2 deletions(-)
diff --git a/python/sedonadb/python/sedonadb/dataframe.py b/python/sedonadb/python/sedonadb/dataframe.py
index a1d0f6c35..d444f88f0 100644
--- a/python/sedonadb/python/sedonadb/dataframe.py
+++ b/python/sedonadb/python/sedonadb/dataframe.py
@@ -1625,6 +1625,66 @@ class DataFrame:
single_file_output,
)
+ def to_csv(
+ self,
+ path: Union[str, Path],
+ *,
+ has_header: bool = True,
+ delimiter: str = ",",
+ ):
+ """Write this DataFrame to CSV
+
+ This is a plain tabular writer. CSV has no geometry representation, so a
+ geometry column (or one nested inside a struct/list) raises an error
+ rather than being written in some opaque encoding; convert it to text
+ first, e.g. `select(ST_AsText(col("geometry")).alias("geometry"))`. To
+ write geometry to a spatial format such as GeoJSON, FlatGeobuf, or
+ GeoPackage, use `to_pyogrio()` instead.
+
+ A path ending in `.csv` is written as a single file; any other path is
+ treated as a directory and written as one CSV file per partition.
+
+ Args:
+ path: A filename or directory to which CSV output should be written.
+ has_header: Whether to write the column names as the first row.
+ delimiter: The single-byte field delimiter to use between values.
+
+ Examples:
+
+ >>> import tempfile
+ >>> sd = sedona.db.connect()
+ >>> td = tempfile.TemporaryDirectory()
+ >>> sd.sql("SELECT 1 AS a, 'x' AS b").to_csv(f"{td.name}/tmp.csv")
+
+ """
+ self._impl.to_csv(str(path), has_header, delimiter)
+
+ def to_json(self, path: Union[str, Path]):
+ """Write this DataFrame to newline-delimited JSON
+
+ This is a plain tabular writer that emits one JSON object per row
+ (NDJSON); it does not produce GeoJSON. Like `to_csv()`, a geometry
+ column (or one nested inside a struct/list) raises an error rather than
+ being written in some opaque encoding; convert it to text first, e.g.
+ `select(ST_AsText(col("geometry")).alias("geometry"))`. To write
+ GeoJSON, use `to_pyogrio()` with a `.geojson` path instead.
+
+ A path ending in `.json` is written as a single file; any other path is
+ treated as a directory and written as one JSON file per partition.
+
+ Args:
+ path: A filename or directory to which JSON output should be written.
+
+ Examples:
+
+ >>> import tempfile
+ >>> sd = sedona.db.connect()
+ >>> td = tempfile.TemporaryDirectory()
+ >>> sd.sql("SELECT 1 AS a, 'x' AS b").to_json(f"{td.name}/tmp.json")
+
+ """
+ self._impl.to_json(str(path))
+
def to_pyogrio(
self,
path: Union[str, Path, io.BytesIO],
diff --git a/python/sedonadb/src/dataframe.rs b/python/sedonadb/src/dataframe.rs
index 07825df45..aca3f78d1 100644
--- a/python/sedonadb/src/dataframe.rs
+++ b/python/sedonadb/src/dataframe.rs
@@ -590,6 +590,39 @@ impl InternalDataFrame {
Ok(())
}
+ fn to_csv<'py>(
+ &self,
+ py: Python<'py>,
+ path: String,
+ has_header: bool,
+ delimiter: &str,
+ ) -> Result<(), PySedonaError> {
+ let bytes = delimiter.as_bytes();
+ if bytes.len() != 1 {
+ return Err(PySedonaError::SedonaPython(format!(
+ "CSV delimiter must be a single byte, got {delimiter:?}"
+ )));
+ }
+
+ wait_for_future(
+ py,
+ &self.runtime,
+ self.inner
+ .clone()
+ .write_sedona_csv(&path, has_header, bytes[0]),
+ )??;
+ Ok(())
+ }
+
+ fn to_json<'py>(&self, py: Python<'py>, path: String) -> Result<(), PySedonaError> {
+ wait_for_future(
+ py,
+ &self.runtime,
+ self.inner.clone().write_sedona_json(&path),
+ )??;
+ Ok(())
+ }
+
fn show<'py>(
&self,
py: Python<'py>,
diff --git a/python/sedonadb/tests/io/test_write_csv_json.py b/python/sedonadb/tests/io/test_write_csv_json.py
new file mode 100644
index 000000000..c92386f47
--- /dev/null
+++ b/python/sedonadb/tests/io/test_write_csv_json.py
@@ -0,0 +1,128 @@
+# 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.
+
+import json
+import tempfile
+from pathlib import Path
+
+import pandas as pd
+import pandas.testing as pdt
+import pytest
+
+from sedonadb._lib import SedonaError
+
+
+def test_to_csv_round_trip(con):
+ with tempfile.TemporaryDirectory() as td:
+ p = Path(td) / "out.csv"
+ con.sql("SELECT 1 AS a, 'x' AS b UNION ALL SELECT 2, 'y'").to_csv(p)
+ # A single file is written when the path ends with ".csv".
+ assert p.is_file()
+ out = con.read.csv(p).sort("a").to_pandas()
+ pdt.assert_frame_equal(out, pd.DataFrame({"a": [1, 2], "b": ["x", "y"]}))
+
+
+def test_to_csv_no_header(con):
+ with tempfile.TemporaryDirectory() as td:
+ p = Path(td) / "out.csv"
+ con.sql("SELECT 1 AS a, 2 AS b").to_csv(p, has_header=False)
+ assert p.read_text() == "1,2\n"
+
+
+def test_to_csv_custom_delimiter(con):
+ with tempfile.TemporaryDirectory() as td:
+ p = Path(td) / "out.csv"
+ con.sql("SELECT 1 AS a, 'x' AS b").to_csv(p, delimiter=";")
+ assert p.read_text() == "a;b\n1;x\n"
+ # And it round-trips when read back with the same delimiter.
+ out = con.read.csv(p, delimiter=";").to_pandas()
+ pdt.assert_frame_equal(out, pd.DataFrame({"a": [1], "b": ["x"]}))
+
+
+def test_to_csv_bad_delimiter_raises(con):
+ with tempfile.TemporaryDirectory() as td:
+ p = Path(td) / "out.csv"
+ with pytest.raises(SedonaError, match="single byte"):
+ con.sql("SELECT 1 AS a").to_csv(p, delimiter=";;")
+
+
+def test_to_csv_directory_output(con):
+ # A path without a ".csv" suffix writes a directory of part file(s).
+ with tempfile.TemporaryDirectory() as td:
+ d = Path(td) / "parts"
+ con.sql("SELECT 1 AS a").to_csv(d)
+ assert d.is_dir()
+ assert list(d.glob("*.csv"))
+
+
+def test_to_csv_geometry_raises(con):
+ # CSV has no geometry representation, so a geometry column is a hard error
+ # (rather than a silent opaque encoding), with a message naming the column.
+ with tempfile.TemporaryDirectory() as td:
+ p = Path(td) / "geo.csv"
+ df = con.sql("SELECT ST_Point(1.0, 2.0) AS geometry, 'a' AS name")
+ with pytest.raises(SedonaError, match='geometry column.*"geometry"'):
+ df.to_csv(p)
+
+
+def test_to_json_geometry_raises(con):
+ with tempfile.TemporaryDirectory() as td:
+ p = Path(td) / "geo.json"
+ df = con.sql("SELECT ST_Point(1.0, 2.0) AS geometry, 'a' AS name")
+ with pytest.raises(SedonaError, match='geometry column.*"geometry"'):
+ df.to_json(p)
+
+
+def test_to_csv_nested_geometry_raises(con):
+ # Geometry nested inside a list/struct (e.g. ST_Dump output) is caught too.
+ with tempfile.TemporaryDirectory() as td:
+ p = Path(td) / "nested.csv"
+ df = con.sql(
+ "SELECT ST_Dump(ST_GeomFromText('MULTIPOINT (0 0, 1 1)')) AS parts"
+ )
+ with pytest.raises(SedonaError, match="geometry column"):
+ df.to_csv(p)
+
+
+def test_to_csv_geometry_as_text_ok(con):
+ # The documented workaround: project geometry to text first.
+ with tempfile.TemporaryDirectory() as td:
+ p = Path(td) / "wkt.csv"
+ con.sql("SELECT ST_AsText(ST_Point(1.0, 2.0)) AS geometry").to_csv(p)
+ assert p.read_text() == "geometry\nPOINT(1 2)\n"
+
+
+def test_to_json_round_trip(con):
+ with tempfile.TemporaryDirectory() as td:
+ p = Path(td) / "out.json"
+ con.sql("SELECT 1 AS a, 'x' AS b UNION ALL SELECT 2, 'y'").to_json(p)
+ assert p.is_file()
+ out = con.read.json(p).sort("a").to_pandas()
+ pdt.assert_frame_equal(out, pd.DataFrame({"a": [1, 2], "b": ["x", "y"]}))
+
+
+def test_to_json_ndjson_format(con):
+ # to_json emits one JSON object per row (NDJSON): one object per line.
+ with tempfile.TemporaryDirectory() as td:
+ p = Path(td) / "out.json"
+ con.sql("SELECT 1 AS a, 'x' AS b UNION ALL SELECT 2, 'y'").to_json(p)
+ lines = p.read_text().splitlines()
+ assert len(lines) == 2
+ assert sorted((json.loads(line) for line in lines), key=lambda r: r["a"]) == [
+ {"a": 1, "b": "x"},
+ {"a": 2, "b": "y"},
+ ]
diff --git a/rust/sedona-schema/src/schema.rs b/rust/sedona-schema/src/schema.rs
index f13dab242..c89f18acd 100644
--- a/rust/sedona-schema/src/schema.rs
+++ b/rust/sedona-schema/src/schema.rs
@@ -17,7 +17,7 @@
use std::collections::HashMap;
-use arrow_schema::Schema;
+use arrow_schema::{DataType, Schema};
use datafusion_common::{DFSchema, Result};
use crate::{datatypes::SedonaType, matchers::ArgMatcher};
@@ -37,6 +37,55 @@ pub trait SedonaSchema {
/// "geography", (3) the column named "geom", (4) the column named "geog",
/// or (5) the first column with a geometry or geography data type.
fn primary_geometry_column_index(&self) -> Result<Option<usize>>;
+
+ /// Return the indices of top-level columns that are, or contain nested, a
+ /// geometry or geography.
+ ///
+ /// Unlike [`Self::geometry_column_indices`], which only matches columns
+ /// whose top-level type is geometry or geography, this descends into
+ /// struct/list/map children so that nested geometry (for example the
+ /// `List<Struct<path, geom>>` produced by `ST_Dump`) is also reported.
+ fn geometry_column_indices_recursive(&self) -> Result<Vec<usize>> {
+ let mut indices = Vec::new();
+ for (i, sedona_type) in self.sedona_types().enumerate() {
+ if sedona_type_contains_geometry(&sedona_type?)? {
+ indices.push(i);
+ }
+ }
+ Ok(indices)
+ }
+}
+
+/// Return whether a parsed [SedonaType] is, or nests, a geometry or geography.
+///
+/// Descends into struct/list/map children, so nested geometry (for example the
+/// `List<Struct<path, geom>>` produced by `ST_Dump`) is detected.
+pub fn sedona_type_contains_geometry(sedona_type: &SedonaType) -> Result<bool> {
+ match sedona_type {
+ SedonaType::Wkb(_, _) | SedonaType::WkbView(_, _) => Ok(true),
+ SedonaType::Arrow(data_type) => data_type_contains_geometry(data_type),
+ _ => Ok(false),
+ }
+}
+
+fn data_type_contains_geometry(data_type: &DataType) -> Result<bool> {
+ match data_type {
+ DataType::Struct(fields) => {
+ for field in fields {
+ if sedona_type_contains_geometry(&SedonaType::from_storage_field(field)?)? {
+ return Ok(true);
+ }
+ }
+ Ok(false)
+ }
+ DataType::List(field) | DataType::LargeList(field) | DataType::FixedSizeList(field, _) => {
+ sedona_type_contains_geometry(&SedonaType::from_storage_field(field)?)
+ }
+ DataType::Map(field, _) => {
+ sedona_type_contains_geometry(&SedonaType::from_storage_field(field)?)
+ }
+ _ => Ok(false),
+ }
}
impl SedonaSchema for DFSchema {
@@ -180,4 +229,41 @@ mod test {
assert_eq!(schema.geometry_column_indices().unwrap(), vec![0]);
assert_eq!(schema.primary_geometry_column_index().unwrap(), Some(0));
}
+
+ #[test]
+ fn geometry_columns_recursive() {
+ // A geometry nested inside a struct/list (e.g. ST_Dump output) is not a
+ // top-level geometry column, but the recursive variant reports it.
+ let nested = Field::new(
+ "parts",
+ DataType::List(
+ Field::new_struct(
+ "item",
+ vec![
+ Field::new("path", DataType::Int32, true),
+ WKB_GEOMETRY.to_storage_field("geom", true).unwrap(),
+ ],
+ true,
+ )
+ .into(),
+ ),
+ true,
+ );
+ let schema = Schema::new(vec![Field::new("id", DataType::Int32, true), nested]);
+ let df_schema: DFSchema = schema.clone().try_into().unwrap();
+
+ // Top-level detection misses the nested geometry; recursive finds it.
+ assert!(schema.geometry_column_indices().unwrap().is_empty());
+ assert_eq!(schema.geometry_column_indices_recursive().unwrap(), vec![1]);
+ assert_eq!(
+ df_schema.geometry_column_indices_recursive().unwrap(),
+ vec![1]
+ );
+
+ // A plain top-level geometry column is still reported.
+ let schema = Schema::new(vec![WKB_GEOMETRY
+ .to_storage_field("geometry", true)
+ .unwrap()]);
+ assert_eq!(schema.geometry_column_indices_recursive().unwrap(), vec![0]);
+ }
}
diff --git a/rust/sedona/src/context.rs b/rust/sedona/src/context.rs
index ae0b69efc..0f4acbd57 100644
--- a/rust/sedona/src/context.rs
+++ b/rust/sedona/src/context.rs
@@ -27,7 +27,7 @@ use crate::{
show::{show_batches, DisplayTableOptions},
};
use arrow_array::RecordBatch;
-use arrow_schema::DataType;
+use arrow_schema::{DataType, Schema};
use async_trait::async_trait;
use datafusion::datasource::file_format::format_as_file_type;
use datafusion::{
@@ -42,6 +42,7 @@ use datafusion::{
sql::parser::{DFParser, Statement},
};
use datafusion::{dataframe::DataFrameWriteOptions, execution::memory_pool::MemoryLimit};
+use datafusion_common::config::{CsvOptions, JsonOptions};
use datafusion_common::not_impl_err;
use datafusion_expr::dml::InsertOp;
use datafusion_expr::sqlparser::dialect::{dialect_from_str, Dialect};
@@ -66,6 +67,7 @@ use sedona_pointcloud::las::{
options::{GeometryEncoding, LasExtraBytes, LasOptions},
};
use sedona_raster_functions::rs_ensure_loaded::RsEnsureLoaded;
+use sedona_schema::schema::SedonaSchema;
#[cfg(feature = "gpu")]
use sedona_spatial_join_gpu::options::GpuOptions;
@@ -682,6 +684,22 @@ pub trait SedonaDataFrame {
options: SedonaWriteOptions,
writer_options: Option<TableGeoParquetOptions>,
) -> Result<Vec<RecordBatch>>;
+
+ /// Write this DataFrame to CSV file(s), rejecting geometry columns
+ ///
+ /// CSV has no geometry representation, so a geometry (or nested geometry)
+ /// column is a hard error with a message pointing at the required
+ /// text projection, rather than silently writing an opaque encoding.
+ async fn write_sedona_csv(
+ self,
+ path: &str,
+ has_header: bool,
+ delimiter: u8,
+ ) -> Result<Vec<RecordBatch>>;
+
+ /// Write this DataFrame to newline-delimited JSON file(s), rejecting
+ /// geometry columns (see [`Self::write_sedona_csv`]).
+ async fn write_sedona_json(self, path: &str) -> Result<Vec<RecordBatch>>;
}
#[async_trait]
@@ -756,6 +774,55 @@ impl SedonaDataFrame for DataFrame {
DataFrame::new(ctx.ctx.state(), plan).collect().await
}
+
+ async fn write_sedona_csv(
+ self,
+ path: &str,
+ has_header: bool,
+ delimiter: u8,
+ ) -> Result<Vec<RecordBatch>, DataFusionError> {
+ reject_geometry_columns(self.schema().as_arrow(), "CSV")?;
+ let csv_options = CsvOptions {
+ has_header: Some(has_header),
+ delimiter,
+ ..Default::default()
+ };
+ self.write_csv(path, DataFrameWriteOptions::new(), Some(csv_options))
+ .await
+ }
+
+ async fn write_sedona_json(self, path: &str) -> Result<Vec<RecordBatch>, DataFusionError> {
+ reject_geometry_columns(self.schema().as_arrow(), "JSON")?;
+ self.write_json(path, DataFrameWriteOptions::new(), None::<JsonOptions>)
+ .await
+ }
+}
+
+/// Reject a schema that contains geometry/geography columns (including columns
+/// nested inside a struct, list, or map) for a text output format that has no
+/// geometry representation, with a message naming the columns and the required
+/// text projection. Lives here (not in the Python binding) so R can reuse it.
+///
+/// The nested-aware geometry detection is provided by
+/// [`SedonaSchema::geometry_column_indices_recursive`].
+fn reject_geometry_columns(schema: &Schema, format: &str) -> Result<()> {
+ let geometry_columns: Vec<&str> = schema
+ .geometry_column_indices_recursive()?
+ .into_iter()
+ .map(|i| schema.field(i).name().as_str())
+ .collect();
+
+ if !geometry_columns.is_empty() {
+ return plan_err!(
+ "Cannot write geometry column(s) {:?} to {}: this format has no geometry \
+ representation. Convert them to text first (e.g. SELECT ST_AsText(geom) AS geom) \
+ before writing.",
+ geometry_columns,
+ format
+ );
+ }
+
+ Ok(())
}
/// A Sedona-specific copy of [DataFrameWriteOptions]