give read_table the same treatment

Author: timsaucer

examples/datafusion-ffi-example/python/tests/_test_table_provider.py

@@ -40,3 +40,7 @@ def test_table_loading():
     ]
 
     assert result == expected
+
+    result = ctx.read_table(table).collect()
+    result = [r.column(0) for r in result]
+    assert result == expected

python/datafusion/context.py

@@ -1177,14 +1177,11 @@ def read_avro(
             self.ctx.read_avro(str(path), schema, file_partition_cols, file_extension)
         )
 
-    def read_table(self, table: Table) -> DataFrame:
-        """Creates a :py:class:`~datafusion.dataframe.DataFrame` from a table.
-
-        For a :py:class:`~datafusion.catalog.Table` such as a
-        :py:class:`~datafusion.catalog.ListingTable`, create a
-        :py:class:`~datafusion.dataframe.DataFrame`.
-        """
-        return DataFrame(self.ctx.read_table(table._inner))
+    def read_table(
+        self, table: Table | TableProviderExportable | DataFrame | pa.dataset.Dataset
+    ) -> DataFrame:
+        """Creates a :py:class:`~datafusion.dataframe.DataFrame` from a table."""
+        return DataFrame(self.ctx.read_table(table))
 
     def execute(self, plan: ExecutionPlan, partitions: int) -> RecordBatchStream:
         """Execute the ``plan`` and return the results."""

python/tests/test_context.py

@@ -312,7 +312,7 @@ def test_register_table(ctx, database):
     assert public.names() == {"csv", "csv1", "csv2", "csv3"}
 
 
-def test_read_table(ctx, database):
+def test_read_table_from_catalog(ctx, database):
     default = ctx.catalog()
     public = default.schema("public")
     assert public.names() == {"csv", "csv1", "csv2"}
@@ -322,6 +322,25 @@ def test_read_table(ctx, database):
     table_df.show()
 
 
+def test_read_table_from_df(ctx):
+    df = ctx.from_pydict({"a": [1, 2]})
+    result = ctx.read_table(df).collect()
+    assert [b.to_pydict() for b in result] == [{"a": [1, 2]}]
+
+
+def test_read_table_from_dataset(ctx):
+    batch = pa.RecordBatch.from_arrays(
+        [pa.array([1, 2, 3]), pa.array([4, 5, 6])],
+        names=["a", "b"],
+    )
+    dataset = ds.dataset([batch])
+
+    result = ctx.read_table(dataset).collect()
+
+    assert result[0].column(0) == pa.array([1, 2, 3])
+    assert result[0].column(1) == pa.array([4, 5, 6])
+
+
 def test_deregister_table(ctx, database):
     default = ctx.catalog()
     public = default.schema("public")

src/context.rs

@@ -1076,7 +1076,8 @@ impl PySessionContext {
         Ok(PyDataFrame::new(df))
     }
 
-    pub fn read_table(&self, table: &PyTable) -> PyDataFusionResult<PyDataFrame> {
+    pub fn read_table(&self, table: Bound<'_, PyAny>) -> PyDataFusionResult<PyDataFrame> {
+        let table = PyTable::new(&table)?;
         let df = self.ctx.read_table(table.table())?;
         Ok(PyDataFrame::new(df))
     }