From 3d4b80e111d6c467518927d32c11c083fc87a443 Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Tue, 15 Oct 2024 12:47:56 -0400 Subject: [PATCH 01/10] API to get schema from the view --- rerun_py/rerun_bindings/rerun_bindings.pyi | 4 ++++ rerun_py/src/dataframe.rs | 17 ++++++++++++++++- 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/rerun_py/rerun_bindings/rerun_bindings.pyi b/rerun_py/rerun_bindings/rerun_bindings.pyi index 0a607b29463c..a21dc8d59384 100644 --- a/rerun_py/rerun_bindings/rerun_bindings.pyi +++ b/rerun_py/rerun_bindings/rerun_bindings.pyi @@ -92,6 +92,10 @@ class RecordingView: increasing when data is sent from a single process. """ + def schema(self) -> Schema: + """The schema of the view.""" + ... + def filter_range_sequence(self, start: int, end: int) -> RecordingView: """ Filter the view to only include data between the given index sequence numbers. diff --git a/rerun_py/src/dataframe.rs b/rerun_py/src/dataframe.rs index 444a01cc79c4..f53ddaf4c2a5 100644 --- a/rerun_py/src/dataframe.rs +++ b/rerun_py/src/dataframe.rs @@ -331,7 +331,6 @@ impl FromPyObject<'_> for ComponentLike { } } -// TODO(jleibs): Maybe this whole thing moves to the View/Recording #[pyclass(frozen, name = "Schema")] #[derive(Clone)] pub struct PySchema { @@ -410,6 +409,22 @@ pub struct PyRecordingView { /// increasing when data is sent from a single process. #[pymethods] impl PyRecordingView { + fn schema(&self, py: Python<'_>) -> PySchema { + let borrowed = self.recording.borrow(py); + let engine = borrowed.engine(); + + let mut query_expression = self.query_expression.clone(); + query_expression.selection = None; + + let query_handle = engine.query(query_expression); + + let contents = query_handle.view_contents(); + + PySchema { + schema: contents.to_vec(), + } + } + #[pyo3(signature = ( *args, columns = None From 5a643f8bc84d7c72126525f89c7af341f1639266 Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Tue, 15 Oct 2024 12:51:53 -0400 Subject: [PATCH 02/10] Test that implicit content filtering works as expected --- rerun_py/tests/unit/test_dataframe.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/rerun_py/tests/unit/test_dataframe.py b/rerun_py/tests/unit/test_dataframe.py index 0f09cdc029d4..e930ed935b70 100644 --- a/rerun_py/tests/unit/test_dataframe.py +++ b/rerun_py/tests/unit/test_dataframe.py @@ -95,7 +95,7 @@ def test_recording_info(self) -> None: assert self.recording.application_id() == APP_ID assert self.recording.recording_id() == str(RECORDING_ID) - def test_schema(self) -> None: + def test_schema_recording(self) -> None: schema = self.recording.schema() assert len(schema.index_columns()) == 3 @@ -112,6 +112,21 @@ def test_schema(self) -> None: assert schema.component_columns()[2].entity_path == "/points" assert schema.component_columns()[2].component_name == "rerun.components.Position3D" + def test_schema_view(self) -> None: + schema = self.recording.view(index="my_index", contents="/**").schema() + + assert len(schema.index_columns()) == 3 + # Position3D, Color + assert len(schema.component_columns()) == 2 + + assert schema.index_columns()[0].name == "log_tick" + assert schema.index_columns()[1].name == "log_time" + assert schema.index_columns()[2].name == "my_index" + assert schema.component_columns()[0].entity_path == "/points" + assert schema.component_columns()[0].component_name == "rerun.components.Color" + assert schema.component_columns()[1].entity_path == "/points" + assert schema.component_columns()[1].component_name == "rerun.components.Position3D" + def test_full_view(self) -> None: view = self.recording.view(index="my_index", contents="points") From a7f01bd6337abe4d02c5ca81f41259bd2ca344df Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Tue, 15 Oct 2024 13:05:55 -0400 Subject: [PATCH 03/10] Make it possible to opt out of hiding indicators and empty columns --- rerun_py/rerun_bindings/rerun_bindings.pyi | 44 +++++++++++++++++++++- rerun_py/src/dataframe.rs | 15 ++++++-- 2 files changed, 54 insertions(+), 5 deletions(-) diff --git a/rerun_py/rerun_bindings/rerun_bindings.pyi b/rerun_py/rerun_bindings/rerun_bindings.pyi index a21dc8d59384..2e1faaaac6a3 100644 --- a/rerun_py/rerun_bindings/rerun_bindings.pyi +++ b/rerun_py/rerun_bindings/rerun_bindings.pyi @@ -166,7 +166,49 @@ class Recording: """A single recording.""" def schema(self) -> Schema: ... - def view(self, *, index: str, contents: ViewContentsLike) -> RecordingView: ... + def view( + self, + *, + index: str, + contents: ViewContentsLike, + include_semantically_empty_columns: bool = False, + include_indicator_columns: bool = False, + include_tombstone_columns: bool = False, + ) -> RecordingView: + """ + Create a view of the recording according to a particular index and content specification. + + Parameters + ---------- + index : str + The index to use for the view. This is typically a timeline name. + contents : ViewContentsLike + The content specification for the view. This must specify one or more EntityPath expressions + as well as optionally a list of columns. If only providing a single EntityPath, use a string, + otherwise provide a dictionary mapping EntityPaths to a list of column names. + include_semantically_empty_columns : bool, optional + Whether to include columns that are semantically empty, by default False. + include_indicator_columns : bool, optional + Whether to include indicator columns, by default False. + include_tombstone_columns : bool, optional + Whether to include tombstone columns, by default False. + Tombstone columns are components used to represent clears. However, even without the clear + tombstone columns, the view will still apply the clear semantics when resolving row contents. + + Examples + -------- + All the data in the recording on the timeline "my_index": + ```python + recording.view(index="my_index", contents="/**") + ``` + + Just the Position3D components in the "points" entity: + ```python + recording.view(index="my_index", contents={"points": "Position3D"}) + ``` + + """ + ... def recording_id(self) -> str: ... def application_id(self) -> str: ... diff --git a/rerun_py/src/dataframe.rs b/rerun_py/src/dataframe.rs index f53ddaf4c2a5..1b874eae2645 100644 --- a/rerun_py/src/dataframe.rs +++ b/rerun_py/src/dataframe.rs @@ -758,15 +758,22 @@ impl PyRecording { } } + #[allow(clippy::fn_params_excessive_bools)] #[pyo3(signature = ( *, index, - contents + contents, + include_semantically_empty_columns = false, + include_indicator_columns = false, + include_tombstone_columns = false, ))] fn view( slf: Bound<'_, Self>, index: &str, contents: Bound<'_, PyAny>, + include_semantically_empty_columns: bool, + include_indicator_columns: bool, + include_tombstone_columns: bool, ) -> PyResult { let borrowed_self = slf.borrow(); @@ -781,9 +788,9 @@ impl PyRecording { let query = QueryExpression { view_contents: Some(contents), - include_semantically_empty_columns: false, - include_indicator_columns: false, - include_tombstone_columns: false, + include_semantically_empty_columns, + include_indicator_columns, + include_tombstone_columns, filtered_index: Some(timeline.timeline), filtered_index_range: None, filtered_index_values: None, From 35e379e43e51257505aeaa752bb2cca5b92eaf60 Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Tue, 15 Oct 2024 13:12:48 -0400 Subject: [PATCH 04/10] Test semantically empty columns --- rerun_py/tests/unit/test_dataframe.py | 21 ++++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/rerun_py/tests/unit/test_dataframe.py b/rerun_py/tests/unit/test_dataframe.py index e930ed935b70..302b50aee8e0 100644 --- a/rerun_py/tests/unit/test_dataframe.py +++ b/rerun_py/tests/unit/test_dataframe.py @@ -49,7 +49,7 @@ def setup_method(self) -> None: rr.init(APP_ID, recording_id=RECORDING_ID) rr.set_time_sequence("my_index", 1) - rr.log("points", rr.Points3D([[1, 2, 3], [4, 5, 6], [7, 8, 9]])) + rr.log("points", rr.Points3D([[1, 2, 3], [4, 5, 6], [7, 8, 9]], radii=[])) rr.set_time_sequence("my_index", 7) rr.log("points", rr.Points3D([[10, 11, 12]], colors=[[255, 0, 0]])) @@ -98,9 +98,10 @@ def test_recording_info(self) -> None: def test_schema_recording(self) -> None: schema = self.recording.schema() + # log_tick, log_time, my_index assert len(schema.index_columns()) == 3 - # Points3DIndicator, Position3D, Color - assert len(schema.component_columns()) == 3 + # Color, Points3DIndicator, Position3D, Radius + assert len(schema.component_columns()) == 4 assert schema.index_columns()[0].name == "log_tick" assert schema.index_columns()[1].name == "log_time" @@ -111,6 +112,8 @@ def test_schema_recording(self) -> None: assert schema.component_columns()[1].component_name == "rerun.components.Points3DIndicator" assert schema.component_columns()[2].entity_path == "/points" assert schema.component_columns()[2].component_name == "rerun.components.Position3D" + assert schema.component_columns()[3].entity_path == "/points" + assert schema.component_columns()[3].component_name == "rerun.components.Radius" def test_schema_view(self) -> None: schema = self.recording.view(index="my_index", contents="/**").schema() @@ -127,6 +130,18 @@ def test_schema_view(self) -> None: assert schema.component_columns()[1].entity_path == "/points" assert schema.component_columns()[1].component_name == "rerun.components.Position3D" + # Force radius to be included + schema = self.recording.view( + index="my_index", + contents="/**", + include_semantically_empty_columns=True, + ).schema() + + assert len(schema.index_columns()) == 3 + # Color, Position3D, Radius + assert len(schema.component_columns()) == 3 + assert schema.component_columns()[2].component_name == "rerun.components.Radius" + def test_full_view(self) -> None: view = self.recording.view(index="my_index", contents="points") From 27f85ce6af05880bb8edf6f3af28ff8330e96d8f Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Tue, 15 Oct 2024 13:29:01 -0400 Subject: [PATCH 05/10] Add accessor for static column --- rerun_py/rerun_bindings/rerun_bindings.pyi | 9 +++++++++ rerun_py/src/dataframe.rs | 5 +++++ 2 files changed, 14 insertions(+) diff --git a/rerun_py/rerun_bindings/rerun_bindings.pyi b/rerun_py/rerun_bindings/rerun_bindings.pyi index 2e1faaaac6a3..96da051309c9 100644 --- a/rerun_py/rerun_bindings/rerun_bindings.pyi +++ b/rerun_py/rerun_bindings/rerun_bindings.pyi @@ -49,6 +49,15 @@ class ComponentColumnDescriptor: """ ... + @property + def is_static(self) -> bool: + """ + Whether the column is static. + + This property is read-only. + """ + ... + class ComponentColumnSelector: """A selector for a component column.""" diff --git a/rerun_py/src/dataframe.rs b/rerun_py/src/dataframe.rs index 1b874eae2645..50fd3cae4650 100644 --- a/rerun_py/src/dataframe.rs +++ b/rerun_py/src/dataframe.rs @@ -125,6 +125,11 @@ impl PyComponentColumnDescriptor { fn component_name(&self) -> &str { &self.0.component_name } + + #[getter] + fn is_static(&self) -> bool { + self.0.is_static + } } impl From for ComponentColumnDescriptor { From a64275564b30caa2950b94ee18a625c8876231db Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Tue, 15 Oct 2024 13:29:22 -0400 Subject: [PATCH 06/10] Test with static column --- rerun_py/tests/unit/test_dataframe.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/rerun_py/tests/unit/test_dataframe.py b/rerun_py/tests/unit/test_dataframe.py index 302b50aee8e0..3cf06956da50 100644 --- a/rerun_py/tests/unit/test_dataframe.py +++ b/rerun_py/tests/unit/test_dataframe.py @@ -52,6 +52,7 @@ def setup_method(self) -> None: rr.log("points", rr.Points3D([[1, 2, 3], [4, 5, 6], [7, 8, 9]], radii=[])) rr.set_time_sequence("my_index", 7) rr.log("points", rr.Points3D([[10, 11, 12]], colors=[[255, 0, 0]])) + rr.log("static_text", rr.TextLog("Hello"), static=True) with tempfile.TemporaryDirectory() as tmpdir: rrd = tmpdir + "/tmp.rrd" @@ -100,23 +101,30 @@ def test_schema_recording(self) -> None: # log_tick, log_time, my_index assert len(schema.index_columns()) == 3 - # Color, Points3DIndicator, Position3D, Radius - assert len(schema.component_columns()) == 4 + # Color, Points3DIndicator, Position3D, Radius, Text, TextIndicator + assert len(schema.component_columns()) == 6 assert schema.index_columns()[0].name == "log_tick" assert schema.index_columns()[1].name == "log_time" assert schema.index_columns()[2].name == "my_index" assert schema.component_columns()[0].entity_path == "/points" assert schema.component_columns()[0].component_name == "rerun.components.Color" + assert schema.component_columns()[0].is_static is False assert schema.component_columns()[1].entity_path == "/points" assert schema.component_columns()[1].component_name == "rerun.components.Points3DIndicator" + assert schema.component_columns()[1].is_static is False assert schema.component_columns()[2].entity_path == "/points" assert schema.component_columns()[2].component_name == "rerun.components.Position3D" + assert schema.component_columns()[2].is_static is False assert schema.component_columns()[3].entity_path == "/points" assert schema.component_columns()[3].component_name == "rerun.components.Radius" + assert schema.component_columns()[3].is_static is False + assert schema.component_columns()[4].entity_path == "/static_text" + assert schema.component_columns()[4].component_name == "rerun.components.Text" + assert schema.component_columns()[4].is_static is True def test_schema_view(self) -> None: - schema = self.recording.view(index="my_index", contents="/**").schema() + schema = self.recording.view(index="my_index", contents="points").schema() assert len(schema.index_columns()) == 3 # Position3D, Color @@ -133,7 +141,7 @@ def test_schema_view(self) -> None: # Force radius to be included schema = self.recording.view( index="my_index", - contents="/**", + contents="points", include_semantically_empty_columns=True, ).schema() @@ -148,7 +156,7 @@ def test_full_view(self) -> None: batches = view.select() table = pa.Table.from_batches(batches, batches.schema) - # my_index, log_time, log_tick, points, colors + # my_index, log_time, log_tick, points, colors, text assert table.num_columns == 5 assert table.num_rows == 2 From a6055d2fb001e7cdf5bfe6e83c039b6fb5512354 Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Tue, 15 Oct 2024 14:09:05 -0400 Subject: [PATCH 07/10] API for selecting static rows only --- crates/store/re_chunk_store/src/dataframe.rs | 8 ++ rerun_py/rerun_bindings/rerun_bindings.pyi | 3 + rerun_py/src/dataframe.rs | 103 ++++++++++++++++--- 3 files changed, 101 insertions(+), 13 deletions(-) diff --git a/crates/store/re_chunk_store/src/dataframe.rs b/crates/store/re_chunk_store/src/dataframe.rs index 9b5cc068144c..64e3d52b5856 100644 --- a/crates/store/re_chunk_store/src/dataframe.rs +++ b/crates/store/re_chunk_store/src/dataframe.rs @@ -62,6 +62,14 @@ impl ColumnDescriptor { Self::Component(descr) => descr.component_name.short_name().to_owned(), } } + + #[inline] + pub fn is_static(&self) -> bool { + match self { + Self::Time(_) => false, + Self::Component(descr) => descr.is_static, + } + } } /// Describes a time column, such as `log_time`. diff --git a/rerun_py/rerun_bindings/rerun_bindings.pyi b/rerun_py/rerun_bindings/rerun_bindings.pyi index 96da051309c9..6f9d5e743a07 100644 --- a/rerun_py/rerun_bindings/rerun_bindings.pyi +++ b/rerun_py/rerun_bindings/rerun_bindings.pyi @@ -170,6 +170,9 @@ class RecordingView: ... def select(self, *args: AnyColumn, columns: Optional[Sequence[AnyColumn]] = None) -> pa.RecordBatchReader: ... + def select_static( + self, *args: AnyColumn, columns: Optional[Sequence[AnyColumn]] = None + ) -> pa.RecordBatchReader: ... class Recording: """A single recording.""" diff --git a/rerun_py/src/dataframe.rs b/rerun_py/src/dataframe.rs index 50fd3cae4650..e219c5ec21fc 100644 --- a/rerun_py/src/dataframe.rs +++ b/rerun_py/src/dataframe.rs @@ -402,6 +402,29 @@ pub struct PyRecordingView { query_expression: QueryExpression, } +impl PyRecordingView { + fn select_args( + args: &Bound<'_, PyTuple>, + columns: Option>, + ) -> PyResult>> { + // Coerce the arguments into a list of `ColumnSelector`s + let args: Vec = args + .iter() + .map(|arg| arg.extract::()) + .collect::>()?; + + if columns.is_some() && !args.is_empty() { + return Err(PyValueError::new_err( + "Cannot specify both `columns` and `args` in `select`.", + )); + } + + let columns = columns.or_else(|| if !args.is_empty() { Some(args) } else { None }); + + Ok(columns.map(|cols| cols.into_iter().map(|col| col.into_selector()).collect())) + } +} + /// A view of a recording restricted to a given index, containing a specific set of entities and components. /// /// Can only be created by calling `view(...)` on a `Recording`. @@ -445,25 +468,79 @@ impl PyRecordingView { let mut query_expression = self.query_expression.clone(); - // Coerce the arguments into a list of `ColumnSelector`s - let args: Vec = args - .iter() - .map(|arg| arg.extract::()) - .collect::>()?; + query_expression.selection = Self::select_args(args, columns)?; - if columns.is_some() && !args.is_empty() { - return Err(PyValueError::new_err( - "Cannot specify both `columns` and `args` in `select`.", - )); - } + let query_handle = engine.query(query_expression); - let columns = columns.or_else(|| if !args.is_empty() { Some(args) } else { None }); + let schema = query_handle.schema(); + let fields: Vec = + schema.fields.iter().map(|f| f.clone().into()).collect(); + let metadata = schema.metadata.clone().into_iter().collect(); + let schema = arrow::datatypes::Schema::new(fields).with_metadata(metadata); - query_expression.selection = - columns.map(|cols| cols.into_iter().map(|col| col.into_selector()).collect()); + // TODO(jleibs): Need to keep the engine alive + /* + let reader = RecordBatchIterator::new( + query_handle + .into_batch_iter() + .map(|batch| batch.try_to_arrow_record_batch()), + std::sync::Arc::new(schema), + ); + */ + let batches = query_handle + .into_batch_iter() + .map(|batch| batch.try_to_arrow_record_batch()) + .collect::>(); + + let reader = RecordBatchIterator::new(batches.into_iter(), std::sync::Arc::new(schema)); + + Ok(PyArrowType(Box::new(reader))) + } + + #[pyo3(signature = ( + *args, + columns = None + ))] + fn select_static( + &self, + py: Python<'_>, + args: &Bound<'_, PyTuple>, + columns: Option>, + ) -> PyResult>> { + let borrowed = self.recording.borrow(py); + let engine = borrowed.engine(); + + let mut query_expression = self.query_expression.clone(); + + // This is a static selection, so we clear the filtered index + query_expression.filtered_index = None; + + // If no columns provided, select all static columns + let static_columns = Self::select_args(args, columns)?.unwrap_or_else(|| { + self.schema(py) + .schema + .iter() + .filter(|col| col.is_static()) + .map(|col| col.clone().into()) + .collect() + }); + + query_expression.selection = Some(static_columns); let query_handle = engine.query(query_expression); + let non_static_cols = query_handle + .selected_contents() + .iter() + .filter(|(_, col)| !col.is_static()) + .collect::>(); + + if !non_static_cols.is_empty() { + return Err(PyValueError::new_err(format!( + "Static selection resulted in non-static columns: {non_static_cols:?}", + ))); + } + let schema = query_handle.schema(); let fields: Vec = schema.fields.iter().map(|f| f.clone().into()).collect(); From b9b4ce880969132b2c2fe805a261c6725a13be95 Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Tue, 15 Oct 2024 14:59:56 -0400 Subject: [PATCH 08/10] Make it possible to iterator all columns in schema --- rerun_py/rerun_bindings/rerun_bindings.pyi | 8 ++++- rerun_py/rerun_bindings/types.py | 3 +- rerun_py/src/dataframe.rs | 40 ++++++++++++++++++++++ 3 files changed, 49 insertions(+), 2 deletions(-) diff --git a/rerun_py/rerun_bindings/rerun_bindings.pyi b/rerun_py/rerun_bindings/rerun_bindings.pyi index 6f9d5e743a07..428287c66f93 100644 --- a/rerun_py/rerun_bindings/rerun_bindings.pyi +++ b/rerun_py/rerun_bindings/rerun_bindings.pyi @@ -1,5 +1,5 @@ import os -from typing import Optional, Sequence +from typing import Iterator, Optional, Sequence, Union import pyarrow as pa @@ -16,6 +16,11 @@ class IndexColumnDescriptor: """ ... + @property + def is_static(self) -> bool: + """ColumnDescriptor interface: always False for Index.""" + ... + class IndexColumnSelector: """A selector for an index column.""" @@ -83,6 +88,7 @@ class ComponentColumnSelector: class Schema: """The schema representing all columns in a [`Recording`][].""" + def __iter__(self) -> Iterator[Union[IndexColumnDescriptor, ComponentColumnDescriptor]]: ... def index_columns(self) -> list[IndexColumnDescriptor]: ... def component_columns(self) -> list[ComponentColumnDescriptor]: ... def column_for(self, entity_path: str, component: ComponentLike) -> Optional[ComponentColumnDescriptor]: ... diff --git a/rerun_py/rerun_bindings/types.py b/rerun_py/rerun_bindings/types.py index 87c219f21c9c..f5b08f84e852 100644 --- a/rerun_py/rerun_bindings/types.py +++ b/rerun_py/rerun_bindings/types.py @@ -12,7 +12,7 @@ from .rerun_bindings import ( ComponentColumnDescriptor as ComponentColumnDescriptor, ComponentColumnSelector as ComponentColumnSelector, - IndexColumnSelector as IndexColumnDescriptor, + IndexColumnDescriptor as IndexColumnDescriptor, IndexColumnSelector as IndexColumnSelector, ) @@ -25,6 +25,7 @@ "IndexColumnSelector", ] + AnyComponentColumn: TypeAlias = Union[ "ComponentColumnDescriptor", "ComponentColumnSelector", diff --git a/rerun_py/src/dataframe.rs b/rerun_py/src/dataframe.rs index e219c5ec21fc..aabcb54ae644 100644 --- a/rerun_py/src/dataframe.rs +++ b/rerun_py/src/dataframe.rs @@ -57,6 +57,12 @@ impl PyIndexColumnDescriptor { fn name(&self) -> &str { self.0.timeline.name() } + + #[allow(clippy::unused_self)] + #[getter] + fn is_static(&self) -> bool { + false + } } impl From for PyIndexColumnDescriptor { @@ -336,6 +342,21 @@ impl FromPyObject<'_> for ComponentLike { } } +#[pyclass] +pub struct SchemaIterator { + iter: std::vec::IntoIter, +} + +#[pymethods] +impl SchemaIterator { + fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> { + slf + } + fn __next__(mut slf: PyRefMut<'_, Self>) -> Option { + slf.iter.next() + } +} + #[pyclass(frozen, name = "Schema")] #[derive(Clone)] pub struct PySchema { @@ -344,6 +365,25 @@ pub struct PySchema { #[pymethods] impl PySchema { + fn __iter__(slf: PyRef<'_, Self>) -> PyResult> { + let py = slf.py(); + let iter = SchemaIterator { + iter: slf + .schema + .clone() + .into_iter() + .map(|col| match col { + ColumnDescriptor::Time(col) => PyIndexColumnDescriptor(col).into_py(py), + ColumnDescriptor::Component(col) => { + PyComponentColumnDescriptor(col).into_py(py) + } + }) + .collect::>() + .into_iter(), + }; + Py::new(slf.py(), iter) + } + fn index_columns(&self) -> Vec { self.schema .iter() From 7cc54ee7b8028af512881aee22d078de8e916354 Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Tue, 15 Oct 2024 15:00:25 -0400 Subject: [PATCH 09/10] Test of select_static --- rerun_py/tests/unit/test_dataframe.py | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/rerun_py/tests/unit/test_dataframe.py b/rerun_py/tests/unit/test_dataframe.py index 3cf06956da50..be1134c96199 100644 --- a/rerun_py/tests/unit/test_dataframe.py +++ b/rerun_py/tests/unit/test_dataframe.py @@ -151,15 +151,28 @@ def test_schema_view(self) -> None: assert schema.component_columns()[2].component_name == "rerun.components.Radius" def test_full_view(self) -> None: - view = self.recording.view(index="my_index", contents="points") + view = self.recording.view(index="my_index", contents="/**") - batches = view.select() - table = pa.Table.from_batches(batches, batches.schema) + table = view.select().read_all() # my_index, log_time, log_tick, points, colors, text + assert table.num_columns == 6 + assert table.num_rows == 2 + + table = view.select( + columns=[col for col in view.schema() if not col.is_static], + ).read_all() + + # my_index, log_time, log_tick, points, colors assert table.num_columns == 5 assert table.num_rows == 2 + table = view.select_static().read_all() + + # text + assert table.num_columns == 1 + assert table.num_rows == 1 + def test_select_columns(self) -> None: view = self.recording.view(index="my_index", contents="points") index_col = rr.dataframe.IndexColumnSelector("my_index") From 50ab88bd1b372b657abc4be8497383cb81a13849 Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Tue, 15 Oct 2024 15:11:36 -0400 Subject: [PATCH 10/10] lint --- rerun_py/src/dataframe.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/rerun_py/src/dataframe.rs b/rerun_py/src/dataframe.rs index aabcb54ae644..d6352aff90f6 100644 --- a/rerun_py/src/dataframe.rs +++ b/rerun_py/src/dataframe.rs @@ -352,6 +352,7 @@ impl SchemaIterator { fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> { slf } + fn __next__(mut slf: PyRefMut<'_, Self>) -> Option { slf.iter.next() }