feat(python/sedonadb): Implement DataFrame.columns (#226)

Author: p-vdp

Cargo.lock

@@ -37,6 +37,7 @@ test = [
     "geoarrow-pyarrow",
     "geopandas",
     "pandas",
+    "polars",
     "pytest",
 ]
 geopandas = [

python/sedonadb/pyproject.toml

@@ -58,6 +58,11 @@ def schema(self):
         """
         return self._impl.schema()
 
+    @property
+    def columns(self) -> list[str]:
+        """Return a list of column names"""
+        return self._impl.columns()
+
     def head(self, n: int = 5) -> "DataFrame":
         """Limit result to the first n rows
 

python/sedonadb/python/sedonadb/dataframe.py

@@ -61,6 +61,16 @@ impl InternalDataFrame {
         PySedonaSchema::new(arrow_schema.clone())
     }
 
+    fn columns(&self) -> Result<Vec<String>, PySedonaError> {
+        Ok(self
+            .inner
+            .schema()
+            .fields()
+            .iter()
+            .map(|f| f.name().to_string())
+            .collect())
+    }
+
     fn primary_geometry_column(&self) -> Result<Option<String>, PySedonaError> {
         Ok(self
             .inner

python/sedonadb/src/dataframe.rs

@@ -162,6 +162,14 @@ def test_schema(con):
         df.schema.field({})
 
 
+def test_columns(con):
+    df = con.sql("SELECT 1 as one, ST_GeomFromWKT('POINT (0 1)') as geom")
+    assert len(df.columns) == 2
+
+    pdf = df.to_pandas()
+    assert set(df.columns) == set(pdf.columns)
+
+
 def test_schema_non_null_crs(con):
     tab = pa.table({"geom": ga.with_crs(ga.as_wkb(["POINT (0 1)"]), gat.OGC_CRS84)})
     df = con.create_data_frame(tab)