paleolimbot commented on code in PR #1044:
URL: https://github.com/apache/sedona-db/pull/1044#discussion_r3653753810
##########
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:
There's a SedonaSchema trait for schema-level geometry inspection where this
could live, too (may have a higher chance of getting reused there)
##########
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:
This is imported in a number of individual tests and can be hoisted to the
module level
--
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]
