Merge pull request #536 from chdb-io/improve/pristine-metadata-optimization

Improve metadata access for pristine DataStore sources

Author: auxten

.gitignore

@@ -264,3 +264,4 @@ libchdb.so
 
 # supposed to be some local rules, not for sharing
 .cursor/rules/
+chdb.code-workspace

datastore/core.py

@@ -684,7 +684,7 @@ def _render_operations(self, operations, start_num=1, verbose=False):
                     lines.append(f"     └─ {k}: {v}")
         return lines
 
-    def explain(self, verbose: bool = False) -> str:
+    def explain(self, verbose: bool = False) -> None:
         """
         Generate and display the execution plan in original operation order.
 
@@ -696,15 +696,12 @@ def explain(self, verbose: bool = False) -> str:
         Args:
             verbose: If True, show additional details like full SQL queries
 
-        Returns:
-            String representation of the execution plan
-
         Example:
             >>> ds = DataStore.from_file("data.csv")
             >>> ds = ds.select('name', 'age').filter(ds.age > 25)
             >>> ds['computed'] = ds['age'] * 2
             >>> ds = ds.filter(ds['age'] < 50)  # Order matters!
-            >>> print(ds.explain())
+            >>> ds.explain()
         """
         from .query_planner import QueryPlanner
 
@@ -867,7 +864,6 @@ def explain(self, verbose: bool = False) -> str:
 
         output = "\n".join(lines)
         print(output)
-        return output
 
     def _is_sql_query(self) -> bool:
         """Check if this DataStore represents a SQL query."""
@@ -1089,6 +1085,58 @@ def _add_lazy_op(self, op: LazyOp):
         self._lazy_ops.append(op)
         self._invalidate_cache()
 
+    def _get_source_df_if_pristine(self) -> "Optional[pd.DataFrame]":
+        """
+        For pristine DataStores (no transformations applied), return the
+        underlying source DataFrame directly. Returns None otherwise.
+
+        "Pristine" means:
+        - Cache is valid (already executed, result available), OR
+        - Exactly one LazyDataFrameSource op with no further ops
+
+        This enables zero-cost metadata access (dtypes, columns, shape)
+        for freshly created or already-executed DataStores.
+        """
+        if self._is_cache_valid():
+            return self._cached_result
+        from .lazy_ops import LazyDataFrameSource
+        if (len(self._lazy_ops) == 1
+                and isinstance(self._lazy_ops[0], LazyDataFrameSource)):
+            return self._lazy_ops[0]._df
+        return None
+
+    def _is_pristine_sql_source(self) -> bool:
+        """Check if this is a raw SQL/file source with no operations."""
+        return (not self._lazy_ops
+                and (self._table_function is not None or self.table_name is not None))
+
+    def _probe_dtypes_from_sql_source(self) -> "pd.Series":
+        """
+        Get pandas dtypes for a pristine SQL source via a LIMIT 0 query.
+
+        Executes SELECT * ... LIMIT 0 which returns an empty DataFrame with
+        correct column types, without transferring any actual data.
+        """
+        if self._executor is None:
+            self.connect()
+
+        if self._table_function:
+            source = self._table_function.to_sql()
+        elif self.table_name:
+            source = format_identifier(self.table_name, self.quote_char)
+        else:
+            return self._get_df().dtypes
+
+        sql = f"SELECT * FROM {source} LIMIT 0"
+        if self._format_settings:
+            parts = []
+            for k, v in self._format_settings.items():
+                parts.append(f"{k}='{v}'" if isinstance(v, str) else f"{k}={v}")
+            sql += " SETTINGS " + ", ".join(parts)
+
+        result = self._executor.execute(sql)
+        return result.to_df().dtypes
+
     def _execute(self):
         """
         Execute all lazy operations into a DataFrame.
@@ -3598,9 +3646,10 @@ def sample(
     @property
     def shape(self):
         """
-        Return the shape (rows, columns) of the query result.
+        Return the shape (rows, columns).
 
-        Works correctly with both SQL queries and executed DataFrames.
+        Optimized for pristine sources: uses source DataFrame directly or
+        SQL COUNT(*) + DESCRIBE, avoiding full data loading.
 
         Returns:
             Tuple of (rows, columns)
@@ -3609,20 +3658,23 @@ def shape(self):
             >>> ds = DataStore.from_file("data.csv")
             >>> rows, cols = ds.select("*").shape
         """
-        # Use _get_df if available (handles caching properly)
+        src = self._get_source_df_if_pristine()
+        if src is not None:
+            return src.shape
+        if self._is_pristine_sql_source():
+            return (self.count_rows(), len(self.columns))
         if hasattr(self, "_get_df"):
-            df = self._get_df()
-        else:
-            df = self.to_df()
-        return df.shape
+            return self._get_df().shape
+        return self.to_df().shape
 
     @property
     def columns(self):
         """
-        Return the column names of the query result.
+        Return the column names.
 
-        Works correctly with both SQL queries and executed DataFrames.
-        Raises DataStoreError if no table is bound (connection-level DataStore).
+        Optimized for pristine sources: reads from the source DataFrame
+        or uses DESCRIBE (no data loading). For complex pipelines,
+        falls back to full execution.
 
         Returns:
             pandas Index of column names
@@ -3632,12 +3684,16 @@ def columns(self):
             >>> cols = ds.select("*").columns
         """
         self._require_table_context("columns")
-        # Use _get_df if available (handles caching properly)
+        src = self._get_source_df_if_pristine()
+        if src is not None:
+            return src.columns
+        if self._is_pristine_sql_source():
+            schema = self.schema()
+            if schema:
+                return pd.Index(schema.keys())
         if hasattr(self, "_get_df"):
-            df = self._get_df()
-        else:
-            df = self.to_df()
-        return df.columns
+            return self._get_df().columns
+        return self.to_df().columns
 
     @columns.setter
     def columns(self, new_columns):

datastore/pandas_compat.py

@@ -348,7 +348,12 @@ def _wrap_result(self, result, operation_name: str = None):
 
     @property
     def dtypes(self):
-        """Return the dtypes in the DataFrame."""
+        """Return the dtypes. Optimized for pristine sources."""
+        src = self._get_source_df_if_pristine()
+        if src is not None:
+            return src.dtypes
+        if self._is_pristine_sql_source():
+            return self._probe_dtypes_from_sql_source()
         return self._get_df().dtypes
 
     @property
@@ -368,17 +373,27 @@ def axes(self):
 
     @property
     def ndim(self):
-        """Return the number of dimensions."""
-        return self._get_df().ndim
+        """Return the number of dimensions (always 2 for DataStore/DataFrame)."""
+        return 2
 
     @property
     def size(self):
-        """Return the number of elements in the DataFrame."""
+        """Return the number of elements. Optimized for pristine sources."""
+        src = self._get_source_df_if_pristine()
+        if src is not None:
+            return src.size
+        if self._is_pristine_sql_source():
+            return self.count_rows() * len(self.columns)
         return self._get_df().size
 
     @property
     def empty(self):
-        """Indicator whether DataFrame is empty."""
+        """Check if DataFrame is empty. Optimized for pristine sources."""
+        src = self._get_source_df_if_pristine()
+        if src is not None:
+            return len(src) == 0
+        if self._is_pristine_sql_source():
+            return self.count_rows() == 0
         return self._get_df().empty
 
     @property

datastore/tests/setup_clickhouse_server.sh

@@ -160,6 +160,7 @@ cat > "${CONFIG_FILE}" << EOF
     </logger>
     <listen_host>127.0.0.1</listen_host>
     <users_config>users.xml</users_config>
+    <max_server_memory_usage_to_ram_ratio>0.25</max_server_memory_usage_to_ram_ratio>
 </clickhouse>
 EOF
 
@@ -170,7 +171,7 @@ cat > "${USERS_FILE}" << EOF
 <clickhouse>
     <profiles>
         <default>
-            <max_memory_usage>10000000000</max_memory_usage>
+            <max_memory_usage>500000000</max_memory_usage>
             <use_uncompressed_cache>0</use_uncompressed_cache>
             <load_balancing>random</load_balancing>
         </default>

datastore/tests/test_cache.py

@@ -5,15 +5,24 @@
 re-execution of the pipeline when repr/__str__ are called multiple times.
 """
 
+import io
 import time
 import unittest
+from contextlib import redirect_stdout
 
 import pandas as pd
 
 from datastore import DataStore, config
 from tests.test_utils import assert_frame_equal
 
 
+def _capture_explain(obj, **kwargs):
+    f = io.StringIO()
+    with redirect_stdout(f):
+        obj.explain(**kwargs)
+    return f.getvalue()
+
+
 class TestCacheBasics(unittest.TestCase):
     """Test basic caching functionality."""
 
@@ -1066,7 +1075,7 @@ def test_explain_shows_cached_source_after_checkpoint(self):
         str(ds)  # checkpoint
 
         # After checkpoint, explain should show DataFrame source
-        explain_output = ds.explain()
+        explain_output = _capture_explain(ds)
         self.assertIn('DataFrame source', explain_output)
 
     def test_explain_shows_new_ops_after_checkpoint(self):
@@ -1079,7 +1088,7 @@ def test_explain_shows_new_ops_after_checkpoint(self):
 
         ds['z'] = ds['y'] + 1  # new op after checkpoint
 
-        explain_output = ds.explain()
+        explain_output = _capture_explain(ds)
 
         # Should show both cached source and new op
         self.assertIn('DataFrame source', explain_output)
@@ -1664,7 +1673,7 @@ def test_explain_shows_correct_engine_indicators(self):
         ds = DataStore.from_dataframe(pd.DataFrame({'a': [1, 2, 3]}))
         ds['b'] = ds['a'] * 2
 
-        explain_output = ds.explain()
+        explain_output = _capture_explain(ds)
 
         # In Pandas mode, explain should show [Pandas] indicators
         self.assertIn('[Pandas]', explain_output,
@@ -1675,7 +1684,7 @@ def test_explain_shows_correct_engine_indicators(self):
         ds2 = DataStore.from_dataframe(pd.DataFrame({'a': [1, 2, 3]}))
         ds2['b'] = ds2['a'] * 2
 
-        explain_output2 = ds2.explain()
+        explain_output2 = _capture_explain(ds2)
 
         # In Auto mode, explain might show either engine
         # The key is that it should have engine indicators

datastore/tests/test_comprehensive_segmented_execution.py

@@ -24,6 +24,7 @@
 import re
 import tempfile
 import unittest
+from contextlib import redirect_stdout
 from datetime import datetime, timedelta
 
 import numpy as np
@@ -40,6 +41,13 @@
 from tests.test_utils import assert_datastore_equals_pandas
 
 
+def _capture_explain(obj, **kwargs):
+    f = io.StringIO()
+    with redirect_stdout(f):
+        obj.explain(**kwargs)
+    return f.getvalue()
+
+
 def verify_segment_engines(lazy_ops, has_sql_source, expected_segments):
     """
     Verify that the execution plan matches expected segment types.
@@ -926,7 +934,7 @@ def test_explain_shows_lazy_ops(self):
         ds['doubled'] = ds['value'].apply(lambda x: x * 2)
         ds = ds.sort_values('doubled')
 
-        explain = ds.explain()
+        explain = _capture_explain(ds)
 
         # === VERIFY HEADER ===
         self.assertIn('Execution Plan (in execution order)', explain, "Should show execution plan header")