jiayuasu commented on code in PR #1044:
URL: https://github.com/apache/sedona-db/pull/1044#discussion_r3654021404
##########
rust/sedona/src/context.rs:
##########
@@ -741,6 +759,80 @@ 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.
+fn reject_geometry_columns(schema: &Schema, format: &str) -> Result<()> {
Review Comment:
Moved it — added `geometry_column_indices_recursive()` as a default method
on the `SedonaSchema` trait (backed by a nested-aware
`sedona_type_contains_geometry` helper), and `reject_geometry_columns` now
calls it. Kept it a default method so existing `SedonaSchema` impls are
unaffected, and added a unit test covering the nested (`List<Struct<...,
geom>>`) case.
##########
python/sedonadb/tests/io/test_write_csv_json.py:
##########
@@ -0,0 +1,134 @@
+# 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
+
+
+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):
+ from sedonadb._lib import SedonaError
Review Comment:
Done — hoisted `from sedonadb._lib import SedonaError` to a module-level
import.
--
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]
