add extra sampling options (#156)

* add options to message model * add options to settings context
MSeal · Jan 6, 2023 · 6f9edd6 · 6f9edd6
1 parent 83690e2
commit 6f9edd6
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 3 deletions.
diff --git a/src/dx/filtering.py b/src/dx/filtering.py
@@ -142,10 +142,17 @@ def handle_resample(
         filters=raw_filters,
     )
 
-    # allow temporary override of the display limit
-    with settings_context(DISPLAY_MAX_ROWS=sample_size):
+    # allow temporary override of the display
+    context_params = dict(
+        DISPLAY_MAX_ROWS=sample_size,
+        DISPLAY_MAX_COLUMNS=msg.num_columns,
+        COLUMN_SAMPLING_METHOD=msg.column_sampling_method,
+        ROW_SAMPLING_METHOD=msg.row_sampling_method,
+    )
+    with settings_context(**context_params):
         logger.debug(
-            f"updating {msg.display_id=} with {min(sample_size, len(resampled_df))}-row resample"
+            f"updating {msg.display_id=} with {min(sample_size, len(resampled_df))}-row resample",
+            **context_params,
         )
         update_display(
             resampled_df,

diff --git a/src/dx/types/filters.py b/src/dx/types/filters.py
@@ -6,6 +6,8 @@
 from pydantic import BaseModel, Field
 from typing_extensions import Annotated
 
+from dx.types.main import DXSamplingMethod
+
 logger = structlog.get_logger(__name__)
 
 
@@ -108,6 +110,9 @@ class DEXResampleMessage(BaseModel):
     filters: List[Annotated[FilterTypes, Field(discriminator="type")]] = Field(default_factory=list)
     limit: int = 50_000
     cell_id: Optional[str] = None
+    num_columns: int = 100
+    column_sampling_method: DXSamplingMethod = DXSamplingMethod.outer
+    row_sampling_method: DXSamplingMethod = DXSamplingMethod.random
 
 
 def clean_pandas_query_column(column: str) -> str: