duckrun 0.2.2__py3-none-any.whl → 0.2.3__py3-none-any.whl

duckrun/stats.py

@@ -0,0 +1,225 @@
+"""
+Delta Lake table statistics functionality for duckrun
+"""
+import duckdb
+from deltalake import DeltaTable
+from datetime import datetime
+
+
+def _table_exists(duckrun_instance, schema_name: str, table_name: str) -> bool:
+    """Check if a specific table exists by trying to query it directly."""
+    try:
+        # For main schema, just use table name directly
+        if schema_name == "main":
+            query = f"SELECT COUNT(*) FROM {table_name} LIMIT 1"
+        else:
+            query = f"SELECT COUNT(*) FROM {schema_name}.{table_name} LIMIT 1"
+        duckrun_instance.con.execute(query)
+        return True
+    except:
+        return False
+
+
+def _schema_exists(duckrun_instance, schema_name: str) -> bool:
+    """Check if a schema exists by trying to show its tables."""
+    try:
+        # For main schema, just show tables
+        if schema_name == "main":
+            query = "SHOW TABLES"
+        else:
+            query = f"SHOW TABLES FROM {schema_name}"
+        duckrun_instance.con.execute(query)
+        return True
+    except:
+        return False
+
+
+def _get_existing_tables_in_schema(duckrun_instance, schema_name: str) -> list:
+    """Get all existing tables in a schema by showing tables, excluding temporary tables."""
+    try:
+        # For main schema, just show tables
+        if schema_name == "main":
+            query = "SHOW TABLES"
+        else:
+            query = f"SHOW TABLES FROM {schema_name}"
+        result = duckrun_instance.con.execute(query).fetchall()
+        if result:
+            # Filter out temporary tables created by stats processing (tbl_0, tbl_1, etc.)
+            tables = [row[0] for row in result]
+            filtered_tables = [tbl for tbl in tables if not tbl.startswith('tbl_')]
+            return filtered_tables
+        return []
+    except:
+        return []
+
+
+def get_stats(duckrun_instance, source: str):
+    """
+    Get comprehensive statistics for Delta Lake tables.
+    
+    Args:
+        duckrun_instance: The Duckrun connection instance
+        source: Can be one of:
+               - Table name: 'table_name' (uses main schema in DuckDB)
+               - Schema.table: 'schema.table_name' (specific table in schema, if multi-schema)
+               - Schema only: 'schema' (all tables in schema, if multi-schema)
+    
+    Returns:
+        Arrow table with statistics including total rows, file count, row groups, 
+        average row group size, file sizes, VORDER status, and timestamp
+    
+    Examples:
+        con = duckrun.connect("tmp/data.lakehouse/test")
+        
+        # Single table in main schema (DuckDB uses 'main', not 'test')
+        stats = con.get_stats('price_today')
+        
+        # Specific table in different schema (only if multi-schema enabled)
+        stats = con.get_stats('aemo.price')
+        
+        # All tables in a schema (only if multi-schema enabled)
+        stats = con.get_stats('aemo')
+    """
+    timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
+    
+    # DuckDB always uses 'main' as the default schema, regardless of connection URL schema
+    duckdb_schema = "main"
+    url_schema = duckrun_instance.schema  # This is from the connection URL path
+    
+    # Parse the source and validate existence
+    if '.' in source:
+        # Format: schema.table - only valid if multi-schema is enabled
+        schema_name, table_name = source.split('.', 1)
+        
+        if not duckrun_instance.scan_all_schemas:
+            raise ValueError(f"Multi-schema format '{source}' not supported. Connection was made to a specific schema '{url_schema}'. Use just the table name '{table_name}' instead.")
+        
+        # Validate the specific table exists in the actual DuckDB schema
+        if not _table_exists(duckrun_instance, schema_name, table_name):
+            raise ValueError(f"Table '{table_name}' does not exist in schema '{schema_name}'")
+        
+        list_tables = [table_name]
+    else:
+        # Could be just table name or schema name
+        if duckrun_instance.scan_all_schemas:
+            # Multi-schema mode: DuckDB has actual schemas
+            # First check if it's a table in main schema
+            if _table_exists(duckrun_instance, duckdb_schema, source):
+                list_tables = [source]
+                schema_name = duckdb_schema
+            # Otherwise, check if it's a schema name
+            elif _schema_exists(duckrun_instance, source):
+                schema_name = source
+                list_tables = _get_existing_tables_in_schema(duckrun_instance, source)
+                if not list_tables:
+                    raise ValueError(f"Schema '{source}' exists but contains no tables")
+            else:
+                raise ValueError(f"Neither table '{source}' in main schema nor schema '{source}' exists")
+        else:
+            # Single-schema mode: tables are in DuckDB's main schema, use URL schema for file paths
+            if _table_exists(duckrun_instance, duckdb_schema, source):
+                # It's a table name
+                list_tables = [source]
+                schema_name = url_schema  # Use URL schema for file path construction
+            elif source == url_schema:
+                # Special case: user asked for stats on the URL schema name - list all tables
+                list_tables = _get_existing_tables_in_schema(duckrun_instance, duckdb_schema)
+                schema_name = url_schema  # Use URL schema for file path construction
+                if not list_tables:
+                    raise ValueError(f"No tables found in schema '{url_schema}'")
+            else:
+                raise ValueError(f"Table '{source}' does not exist in the current context (schema: {url_schema})")
+    
+    # Use the existing connection
+    con = duckrun_instance.con
+    
+    print(f"Processing {len(list_tables)} tables: {list_tables}")
+    
+    for idx, tbl in enumerate(list_tables):
+        # Construct lakehouse path using ABFSS URL
+        table_path = f"abfss://{duckrun_instance.workspace}@{duckrun_instance.storage_account}.dfs.fabric.microsoft.com/{duckrun_instance.lakehouse_name}.Lakehouse/Tables/{schema_name}/{tbl}"
+        
+        try:
+            dt = DeltaTable(table_path)
+            xx = dt.get_add_actions(flatten=True).to_pydict()
+            
+            # Check if VORDER exists
+            vorder = 'tags.VORDER' in xx.keys()
+            
+            # Calculate total size
+            total_size = sum(xx['size_bytes']) if xx['size_bytes'] else 0
+            
+            # Get Delta files
+            delta_files = dt.files()
+            delta = [table_path + "/" + f for f in delta_files]
+            
+            # Check if table has any files
+            if not delta:
+                # Empty table - create empty temp table
+                con.execute(f'''
+                    CREATE OR REPLACE TEMP TABLE tbl_{idx} AS
+                    SELECT 
+                        '{tbl}' as tbl,
+                        'empty' as file_name,
+                        0 as num_rows,
+                        0 as num_row_groups,
+                        0 as size,
+                        {vorder} as vorder,
+                        '{timestamp}' as timestamp
+                    WHERE false
+                ''')
+            else:
+                # Get parquet metadata and create temp table
+                con.execute(f'''
+                    CREATE OR REPLACE TEMP TABLE tbl_{idx} AS
+                    SELECT 
+                        '{tbl}' as tbl,
+                        file_name,
+                        num_rows,
+                        num_row_groups,
+                        CEIL({total_size}/(1024*1024)) as size,
+                        {vorder} as vorder,
+                        '{timestamp}' as timestamp
+                    FROM parquet_file_metadata({delta})
+                ''')
+            
+        except Exception as e:
+            print(f"Warning: Could not process table '{tbl}': {e}")
+            # Create empty temp table for failed tables
+            con.execute(f'''
+                CREATE OR REPLACE TEMP TABLE tbl_{idx} AS
+                SELECT 
+                    '{tbl}' as tbl,
+                    'error' as file_name,
+                    0 as num_rows,
+                    0 as num_row_groups,
+                    0 as size,
+                    false as vorder,
+                    '{timestamp}' as timestamp
+                WHERE false
+            ''')
+    
+    # Union all temp tables
+    union_parts = [f'SELECT * FROM tbl_{i}' for i in range(len(list_tables))]
+    union_query = ' UNION ALL '.join(union_parts)
+    
+    # Generate final summary
+    final_result = con.execute(f'''
+        SELECT 
+            tbl,
+            SUM(num_rows) as total_rows,
+            COUNT(*) as num_files,
+            SUM(num_row_groups) as num_row_group,
+            CAST(CEIL(SUM(num_rows)::DOUBLE / NULLIF(SUM(num_row_groups), 0)) AS INTEGER) as average_row_group,
+            MIN(size) as file_size_MB,
+            ANY_VALUE(vorder) as vorder,
+            ANY_VALUE(timestamp) as timestamp
+        FROM ({union_query})
+        WHERE tbl IS NOT NULL
+        GROUP BY tbl
+        ORDER BY total_rows DESC
+    ''').fetch_arrow_table()
+    
+    return final_result
+
+

duckrun/writer.py

@@ -0,0 +1,165 @@
+"""
+Delta Lake writer functionality for duckrun - Spark-style write API
+"""
+from deltalake import DeltaTable, write_deltalake
+
+
+# Row Group configuration for optimal Delta Lake performance
+RG = 8_000_000
+
+
+def _build_write_deltalake_args(path, df, mode, schema_mode=None, partition_by=None):
+    """
+    Build arguments for write_deltalake based on requirements:
+    - If schema_mode='merge': use rust engine (no row group params)
+    - Otherwise: use pyarrow engine with row group optimization
+    """
+    args = {
+        'table_or_uri': path,
+        'data': df,
+        'mode': mode
+    }
+    
+    # Add partition_by if specified
+    if partition_by:
+        args['partition_by'] = partition_by
+    
+    # Engine selection based on schema_mode
+    if schema_mode == 'merge':
+        # Use rust engine for schema merging (no row group params supported)
+        args['schema_mode'] = 'merge'
+        args['engine'] = 'rust'
+    else:
+        # Use pyarrow engine with row group optimization (default)
+        args['max_rows_per_file'] = RG
+        args['max_rows_per_group'] = RG
+        args['min_rows_per_group'] = RG
+    
+    return args
+
+
+class DeltaWriter:
+    """Spark-style write API for Delta Lake"""
+    
+    def __init__(self, relation, duckrun_instance):
+        self.relation = relation
+        self.duckrun = duckrun_instance
+        self._format = "delta"
+        self._mode = "overwrite"
+        self._schema_mode = None
+        self._partition_by = None
+    
+    def format(self, format_type: str):
+        """Set output format (only 'delta' supported)"""
+        if format_type.lower() != "delta":
+            raise ValueError(f"Only 'delta' format is supported, got '{format_type}'")
+        self._format = "delta"
+        return self
+    
+    def mode(self, write_mode: str):
+        """Set write mode: 'overwrite' or 'append'"""
+        if write_mode not in {"overwrite", "append"}:
+            raise ValueError(f"Mode must be 'overwrite' or 'append', got '{write_mode}'")
+        self._mode = write_mode
+        return self
+    
+    def option(self, key: str, value):
+        """Set write option (Spark-compatible)"""
+        if key == "mergeSchema":
+            if str(value).lower() in ("true", "1"):
+                self._schema_mode = "merge"
+            else:
+                self._schema_mode = None
+        else:
+            raise ValueError(f"Unsupported option: {key}")
+        return self
+    
+    def partitionBy(self, *columns):
+        """Set partition columns (Spark-compatible)"""
+        if len(columns) == 1 and isinstance(columns[0], (list, tuple)):
+            # Handle partitionBy(["col1", "col2"]) case
+            self._partition_by = list(columns[0])
+        else:
+            # Handle partitionBy("col1", "col2") case
+            self._partition_by = list(columns)
+        return self
+    
+    def saveAsTable(self, table_name: str):
+        """Save query result as Delta table"""
+        if self._format != "delta":
+            raise RuntimeError(f"Only 'delta' format is supported, got '{self._format}'")
+        
+        if "." in table_name:
+            schema, table = table_name.split(".", 1)
+        else:
+            schema = self.duckrun.schema
+            table = table_name
+        
+        self.duckrun._create_onelake_secret()
+        path = f"{self.duckrun.table_base_url}{schema}/{table}"
+        df = self.relation.record_batch()
+        
+        # Build write arguments based on schema_mode and partition_by
+        write_args = _build_write_deltalake_args(
+            path, df, self._mode, 
+            schema_mode=self._schema_mode,
+            partition_by=self._partition_by
+        )
+        
+        engine_info = f" (engine=rust, schema_mode=merge)" if self._schema_mode == 'merge' else " (engine=pyarrow)"
+        partition_info = f" partitioned by {self._partition_by}" if self._partition_by else ""
+        print(f"Writing to Delta table: {schema}.{table} (mode={self._mode}){engine_info}{partition_info}")
+        
+        write_deltalake(**write_args)
+        
+        # Create view with appropriate schema qualification
+        # If user explicitly specified schema.table, create view with schema qualification
+        # If user just specified table, create view in current schema
+        if "." in table_name:
+            # User explicitly specified schema.table - create qualified view
+            view_name = f"{schema}.{table}"
+            # Ensure the schema exists before creating the view
+            self.duckrun.con.sql(f"CREATE SCHEMA IF NOT EXISTS {schema}")
+        else:
+            # User specified just table name - create view in current schema
+            view_name = table
+            
+        self.duckrun.con.sql(f"DROP VIEW IF EXISTS {view_name}")
+        self.duckrun.con.sql(f"""
+            CREATE OR REPLACE VIEW {view_name}
+            AS SELECT * FROM delta_scan('{path}')
+        """)
+        
+        dt = DeltaTable(path)
+        
+        if self._mode == "overwrite":
+            dt.vacuum(retention_hours=0, dry_run=False, enforce_retention_duration=False)
+            dt.cleanup_metadata()
+            print(f"✅ Table {schema}.{table} created/overwritten")
+        else:
+            file_count = len(dt.file_uris())
+            if file_count > self.duckrun.compaction_threshold:
+                print(f"Compacting {schema}.{table} ({file_count} files)")
+                dt.optimize.compact()
+                dt.vacuum(dry_run=False)
+                dt.cleanup_metadata()
+            print(f"✅ Data appended to {schema}.{table}")
+        
+        return table
+
+
+class QueryResult:
+    """Wrapper for DuckDB relation with write API"""
+    
+    def __init__(self, relation, duckrun_instance):
+        self.relation = relation
+        self.duckrun = duckrun_instance
+    
+    @property
+    def write(self):
+        """Access write API"""
+        return DeltaWriter(self.relation, self.duckrun)
+    
+    def __getattr__(self, name):
+        """Delegate all other methods to underlying DuckDB relation"""
+        return getattr(self.relation, name)

{duckrun-0.2.2.dist-info → duckrun-0.2.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: duckrun
-Version: 0.2.2
+Version: 0.2.3
 Summary: Lakehouse task runner powered by DuckDB for Microsoft Fabric
 Author: mim
 License: MIT

duckrun-0.2.3.dist-info/RECORD

@@ -0,0 +1,11 @@
+duckrun/__init__.py,sha256=L0jRtD9Ld8Ti4e6GRvPDdHvkQCFAPHM43GSP7ARh6EM,241
+duckrun/core.py,sha256=LN5rc5B3HLimgslZdC8tLKe3rjTl_KD8WxCh1qoJhdM,16443
+duckrun/files.py,sha256=xba0juMEQPgaznDudmXcwaGH0wv-6aCoHmV_cNF6Y7I,10665
+duckrun/runner.py,sha256=X5g-57OCHQZ7USKpcBbhYGUcZwLQny2x147DLKrV32c,11417
+duckrun/stats.py,sha256=jLEkxNo7MjibPMpjMsXyedrJqv9-BAnP1C0L2a7H8Z8,9417
+duckrun/writer.py,sha256=eWrGtDQTbXi8H3sSt2WucYTdEQUjK97KmQxzCbqAuMs,6221
+duckrun-0.2.3.dist-info/licenses/LICENSE,sha256=-DeQQwdbCbkB4507ZF3QbocysB-EIjDtaLexvqRkGZc,1083
+duckrun-0.2.3.dist-info/METADATA,sha256=CpJvtR9l8c9b1AV9-KnjN4fZODE_3oJxS3omz4p-qlc,18339
+duckrun-0.2.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+duckrun-0.2.3.dist-info/top_level.txt,sha256=BknMEwebbUHrVAp3SC92ps8MPhK7XSYsaogTvi_DmEU,8
+duckrun-0.2.3.dist-info/RECORD,,

duckrun-0.2.2.dist-info/RECORD

@@ -1,7 +0,0 @@
-duckrun/__init__.py,sha256=L0jRtD9Ld8Ti4e6GRvPDdHvkQCFAPHM43GSP7ARh6EM,241
-duckrun/core.py,sha256=VqfTL4fFE-XUXXsDy9VRFEPSQ21dfrkCGH_06C9CLNg,39416
-duckrun-0.2.2.dist-info/licenses/LICENSE,sha256=-DeQQwdbCbkB4507ZF3QbocysB-EIjDtaLexvqRkGZc,1083
-duckrun-0.2.2.dist-info/METADATA,sha256=J_Vw7Ps5afPRkofvyo-r7wufizjS431XgXpHdwaKwyo,18339
-duckrun-0.2.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-duckrun-0.2.2.dist-info/top_level.txt,sha256=BknMEwebbUHrVAp3SC92ps8MPhK7XSYsaogTvi_DmEU,8
-duckrun-0.2.2.dist-info/RECORD,,